Fix failing tests due to iast collector not initialized

DataDog · Mar 14, 2024 · 02bd708 · 02bd708
1 parent 137b646
commit 02bd708
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 5 deletions.
diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastSystem.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastSystem.java
@@ -82,7 +82,7 @@ public static void start(
     VERBOSITY = config.getIastTelemetryVerbosity();
     LOGGER.debug("IAST is starting: debug={}, verbosity={}", DEBUG, VERBOSITY);
     if (VERBOSITY != Verbosity.OFF) {
-      IastMetricCollector.init(new IastMetricCollector());
+      IastMetricCollector.register(new IastMetricCollector());
     }
     final Reporter reporter = new Reporter(config, AgentTaskScheduler.INSTANCE);
     final boolean globalContext = config.getIastContextMode() == GLOBAL;

diff --git a/...ling/src/test/groovy/datadog/trace/agent/tooling/iast/IastPostProcessorFactoryTest.groovy b/...ling/src/test/groovy/datadog/trace/agent/tooling/iast/IastPostProcessorFactoryTest.groovy
@@ -12,9 +12,13 @@ import net.bytebuddy.implementation.bytecode.assign.Assigner
 import net.bytebuddy.jar.asm.MethodVisitor
 import net.bytebuddy.jar.asm.Opcodes
 import net.bytebuddy.jar.asm.Type
+import spock.lang.Shared
 
 class IastPostProcessorFactoryTest extends DDSpecification {
 
+  @Shared
+  protected static final IastMetricCollector ORIGINAL_COLLECTOR = IastMetricCollector.INSTANCE
+
   private static final Type COLLECTOR_TYPE = Type.getType(IastMetricCollector)
   private static final Type METRIC_TYPE = Type.getType(IastMetric)
 
@@ -23,6 +27,14 @@ class IastPostProcessorFactoryTest extends DDSpecification {
     static void exit() {}
   }
 
+  void setup() {
+    IastMetricCollector.register(new IastMetricCollector())
+  }
+
+  void cleanup() {
+    IastMetricCollector.register(ORIGINAL_COLLECTOR)
+  }
+
   void 'test factory for non annotated'() {
     given:
     final method = new MethodDescription.ForLoadedMethod(NonAnnotatedAdvice.getDeclaredMethod('exit'))

diff --git a/internal-api/src/main/java/datadog/trace/api/iast/telemetry/IastMetricCollector.java b/internal-api/src/main/java/datadog/trace/api/iast/telemetry/IastMetricCollector.java
@@ -32,7 +32,7 @@ public class IastMetricCollector implements MetricCollector<IastMetricCollector.
 
   private static IastMetricCollector INSTANCE = new NoOpInstance();
 
-  public static void init(final IastMetricCollector collector) {
+  public static void register(final IastMetricCollector collector) {
     INSTANCE = collector;
   }
 

diff --git a/internal-api/src/test/groovy/datadog/trace/api/iast/telemetry/IastMetricCollectorTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/iast/telemetry/IastMetricCollectorTest.groovy
@@ -24,6 +24,9 @@ class IastMetricCollectorTest extends DDSpecification {
   @Shared
   protected static final AgentTracer.TracerAPI ORIGINAL_TRACER = AgentTracer.get()
 
+  @Shared
+  protected static final IastMetricCollector ORIGINAL_COLLECTOR = IastMetricCollector.INSTANCE
+
   private HasMetricCollector iastCtx
   private RequestContext ctx
   private AgentSpan span
@@ -48,6 +51,7 @@ class IastMetricCollectorTest extends DDSpecification {
   void cleanup() {
     executor.shutdown()
     AgentTracer.forceRegister(ORIGINAL_TRACER)
+    IastMetricCollector.register(ORIGINAL_COLLECTOR)
   }
 
   void 'test empty collector'() {
@@ -68,6 +72,8 @@ class IastMetricCollectorTest extends DDSpecification {
     final value = 5
     final total = times * value
     final latch = new CountDownLatch(1)
+    final globalCollector = new IastMetricCollector()
+    IastMetricCollector.register(globalCollector)
     final requestCollector = new IastMetricCollector()
     final random = new Random()
     final metrics = IastMetric.values()
@@ -89,11 +95,11 @@ class IastMetricCollectorTest extends DDSpecification {
     latch.countDown()
     futures*.get(10, TimeUnit.SECONDS).size()
     requestCollector.prepareMetrics()
-    IastMetricCollector.get().merge(requestCollector.drain())
+    globalCollector.merge(requestCollector.drain())
 
     then:
-    IastMetricCollector.get().prepareMetrics()
-    final result = IastMetricCollector.get().drain()
+    globalCollector.prepareMetrics()
+    final result = globalCollector.drain()
     final computedTotal = result*.value.sum() as long
     computedTotal == total
   }

diff --git a/telemetry/src/test/groovy/datadog/telemetry/metric/IastMetricPeriodicActionTest.groovy b/telemetry/src/test/groovy/datadog/telemetry/metric/IastMetricPeriodicActionTest.groovy
@@ -6,11 +6,23 @@ import datadog.trace.api.iast.SourceTypes
 import datadog.trace.api.iast.telemetry.IastMetric
 import datadog.trace.api.iast.telemetry.IastMetricCollector
 import groovy.transform.CompileDynamic
+import spock.lang.Shared
 import spock.lang.Specification
 
 @CompileDynamic
 class IastMetricPeriodicActionTest extends Specification {
 
+  @Shared
+  protected static final IastMetricCollector ORIGINAL_COLLECTOR = IastMetricCollector.INSTANCE
+
+  void setup() {
+    IastMetricCollector.register(new IastMetricCollector())
+  }
+
+  void cleanup() {
+    IastMetricCollector.register(ORIGINAL_COLLECTOR)
+  }
+
   void 'test metric'() {
     given:
     final action = new IastMetricPeriodicAction()