From f34dfb52a1de9aa1ba4766a2b8063e18169f269d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?=
 <manuel-alvarez-alvarez@users.noreply.github.com>
Date: Tue, 23 Apr 2024 10:41:48 +0200
Subject: [PATCH] Fix TelemetryRequestEndedHandlerTest (#6934)

---
 .../TelemetryRequestEndedHandlerTest.groovy   | 22 ++++++++++++-------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/telemetry/TelemetryRequestEndedHandlerTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/telemetry/TelemetryRequestEndedHandlerTest.groovy
index 639a01a9275..8969de34e6d 100644
--- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/telemetry/TelemetryRequestEndedHandlerTest.groovy
+++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/telemetry/TelemetryRequestEndedHandlerTest.groovy
@@ -12,8 +12,8 @@ import datadog.trace.api.iast.telemetry.IastMetric
 import datadog.trace.api.iast.telemetry.IastMetricCollector
 import datadog.trace.api.iast.telemetry.Verbosity
 import datadog.trace.api.internal.TraceSegment
-import groovy.transform.CompileDynamic
 import groovy.transform.ToString
+import spock.lang.Shared
 
 import static datadog.trace.api.iast.telemetry.IastMetric.EXECUTED_SINK
 import static datadog.trace.api.iast.telemetry.IastMetric.EXECUTED_SOURCE
@@ -23,18 +23,24 @@ import static datadog.trace.api.iast.telemetry.IastMetric.REQUEST_TAINTED
 import static datadog.trace.api.iast.telemetry.IastMetric.Scope
 import static datadog.trace.api.iast.telemetry.IastMetric.TRACE_METRIC_PREFIX
 
-@CompileDynamic
 class TelemetryRequestEndedHandlerTest extends IastModuleImplTestBase {
 
+  @Shared
+  final IastMetricCollector defaultCollector = IastMetricCollector.get()
+
   protected RequestEndedHandler delegate
   protected IastMetricCollector globalCollector
 
+  @Override
+  void cleanupSpec() {
+    IastMetricCollector.register(defaultCollector)
+  }
+
   void setup() {
     InstrumentationBridge.clearIastModules()
     delegate = Spy(new RequestEndedHandler(dependencies))
-    globalCollector = IastMetricCollector.get()
-    globalCollector.prepareMetrics()
-    globalCollector.drain()
+    globalCollector = new IastMetricCollector()
+    IastMetricCollector.register(globalCollector)
     ctx.taintedObjects = TaintedObjectsWithTelemetry.build(Verbosity.DEBUG, ctx)
     ctx.collector = new IastMetricCollector()
   }
@@ -116,18 +122,18 @@ class TelemetryRequestEndedHandlerTest extends IastModuleImplTestBase {
     }
 
     where:
-    metrics                                                                       | description
+    metrics                                                                | description
     [
       metric(REQUEST_TAINTED, 123),
       metric(EXECUTED_SOURCE, SourceTypes.REQUEST_PARAMETER_VALUE, 2),
       metric(EXECUTED_SOURCE, SourceTypes.REQUEST_HEADER_VALUE, 4),
       metric(EXECUTED_SINK, VulnerabilityTypes.SQL_INJECTION, 1),
       metric(EXECUTED_SINK, VulnerabilityTypes.COMMAND_INJECTION, 2),
-    ]                                                                             | 'List of only request scoped metrics'
+    ]                                                                      | 'List of only request scoped metrics'
     [
       metric(REQUEST_TAINTED, 123),
       metric(INSTRUMENTED_SOURCE, SourceTypes.REQUEST_PARAMETER_VALUE, 2),
-    ]                                                                             | 'Mix between global and request scoped metrics'
+    ]                                                                      | 'Mix between global and request scoped metrics'
   }
 
   private static String getSpanTagValue(final IastMetric metric, final Byte tagValue = null) {