Add API to capture logs in unit tests

dd-java-agent/agent-iast/build.gradle

@@ -62,6 +62,8 @@ dependencies {
   testImplementation group: 'io.grpc', name: 'grpc-core', version: grpcVersion
   testImplementation group: 'io.grpc', name: 'grpc-protobuf', version: grpcVersion
 
+  testImplementation libs.logback.classic
+
   jmh project(':utils:test-utils')
   jmh project(':dd-trace-core')
   jmh project(':dd-java-agent:agent-builder')

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/IastSystemTest.groovy

@@ -21,6 +21,7 @@ import datadog.trace.api.iast.sink.StacktraceLeakModule
 import datadog.trace.api.iast.sink.XContentTypeModule
 import datadog.trace.api.internal.TraceSegment
 import datadog.trace.bootstrap.Agent
+import datadog.trace.test.logging.TestLogCollector
 import datadog.trace.test.util.DDSpecification
 
 import static com.datadog.iast.test.TaintedObjectsUtils.noOpTaintedObjects
@@ -33,8 +34,13 @@ class IastSystemTest extends DDSpecification {
     InstrumentationBridge.clearIastModules()
   }
 
+  def cleanup() {
+    TestLogCollector.disable()
+  }
+
   void 'start'() {
     given:
+    TestLogCollector.enable()
     final ig = new InstrumentationGateway()
     final ss = Spy(ig.getSubscriptionService(RequestContextSlot.IAST))
     final cbp = ig.getCallbackProvider(RequestContextSlot.IAST)
@@ -47,7 +53,6 @@ class IastSystemTest extends DDSpecification {
     }
     final igSpanInfo = Mock(IGSpanInfo)
 
-
     when:
     IastSystem.start(ss)
 
@@ -57,6 +62,7 @@ class IastSystemTest extends DDSpecification {
     1 * ss.registerCallback(Events.get().requestHeader(), _)
     1 * ss.registerCallback(Events.get().grpcServerRequestMessage(), _)
     0 * _
+    TestLogCollector.drainCapturedLogs().any { it.message.contains('IAST is starting') }
 
     when:
     final startCallback = cbp.getCallback(Events.get().requestStarted())

dd-java-agent/agent-iast/src/test/resources/logback-test.xml

@@ -0,0 +1,12 @@
+<configuration>
+  <appender name="TEST" class="datadog.trace.test.logging.TestLogbackAppender"/>
+  <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+    <encoder>
+      <pattern>%d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n</pattern>
+    </encoder>
+  </appender>
+  <root level="DEBUG">
+    <appender-ref ref="TEST"/>
+    <appender-ref ref="STDOUT"/>
+  </root>
+</configuration>

utils/test-utils/build.gradle

@@ -9,4 +9,7 @@ dependencies {
 
   api group: 'com.github.stefanbirkner', name: 'system-rules', version: '1.19.0'
   api group: 'commons-fileupload', name: 'commons-fileupload', version: '1.5'
+
+  compileOnly libs.logback.core
+  compileOnly libs.logback.classic
 }

utils/test-utils/src/main/java/datadog/trace/test/logging/CapturedLog.java

@@ -0,0 +1,50 @@
+package datadog.trace.test.logging;
+
+import java.util.Arrays;
+import java.util.Objects;
+import org.slf4j.Marker;
+
+public final class CapturedLog {
+  public final Marker marker;
+  public final String level;
+  public final String template;
+  public final Object[] arguments;
+  public final String message;
+
+  public CapturedLog(
+      final Marker marker,
+      final String level,
+      final String template,
+      final Object[] arguments,
+      final String message) {
+    this.marker = marker;
+    this.level = level;
+    this.template = template;
+    this.arguments = arguments;
+    this.message = message;
+  }
+
+  @Override
+  public String toString() {
+    return "CapturedLog{"
+        + "marker="
+        + marker
+        + ", level='"
+        + level
+        + '\''
+        + ", template='"
+        + template
+        + '\''
+        + ", arguments=..."
+        + ", message='"
+        + message
+        + '\''
+        + '}';
+  }
+
+  @Override
+  public int hashCode() {
+    final int argsHash = arguments == null ? 0 : Arrays.hashCode(arguments);
+    return Objects.hash(marker, level, template, argsHash, message);
+  }
+}

utils/test-utils/src/main/java/datadog/trace/test/logging/TestLogCollector.java

@@ -0,0 +1,68 @@
+package datadog.trace.test.logging;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.atomic.AtomicReference;
+
+public final class TestLogCollector {
+  static final TestLogCollector INSTANCE = new TestLogCollector();
+
+  private final AtomicReference<LinkedBlockingDeque<CapturedLog>> logs =
+      new AtomicReference<>(null);
+
+  private TestLogCollector() {}
+
+  /**
+   * Enable the test log collector. This should be called before any test that needs the logs.
+   * {@link #disable()} must always be called at cleanup.
+   */
+  public static void enable() {
+    INSTANCE._enable();
+  }
+
+  /** Must be called at least once after {@link #enable()} to cleanup the test log collector. */
+  public static void disable() {
+    INSTANCE._disable();
+  }
+
+  /**
+   * Get all captured logs and clear the internal buffer. {@link #enable()} must have been called
+   * before.
+   */
+  public static List<CapturedLog> drainCapturedLogs() {
+    return INSTANCE._drainCapturedLogs();
+  }
+
+  private void _enable() {
+    final LinkedBlockingDeque<CapturedLog> logs = new LinkedBlockingDeque<>();
+    if (!this.logs.compareAndSet(null, logs)) {
+      throw new IllegalStateException("TestLogCollector was enabled without prior cleanup");
+    }
+  }
+
+  private List<CapturedLog> _drainCapturedLogs() {
+    final List<CapturedLog> result = new ArrayList<>();
+    final LinkedBlockingDeque<CapturedLog> logs = this.logs.get();
+    if (logs == null) {
+      throw new IllegalStateException("TestLogCollector was not enabled before draining logs");
+    }
+    logs.drainTo(result);
+    return result;
+  }
+
+  private void _disable() {
+    final LinkedBlockingDeque<CapturedLog> logs = this.logs.getAndSet(null);
+    if (logs == null) {
+      return;
+    }
+    logs.clear();
+  }
+
+  void addLog(final CapturedLog log) {
+    final LinkedBlockingDeque<CapturedLog> logs = this.logs.get();
+    if (logs != null) {
+      logs.add(log);
+    }
+  }
+}

utils/test-utils/src/main/java/datadog/trace/test/logging/TestLogbackAppender.java

@@ -0,0 +1,39 @@
+package datadog.trace.test.logging;
+
+import ch.qos.logback.classic.spi.ILoggingEvent;
+import ch.qos.logback.core.AppenderBase;
+
+/**
+ * Logback appender that captures logs for testing.
+ *
+ * <p>To set this up, add the following to your logback-test.xml:
+ *
+ * <pre>{@code
+ * <configuration>
+ *   <appender name="TEST" class="datadog.trace.test.logging.TestLogbackAppender" />
+ *   <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+ *     <encoder>
+ *       <pattern>%d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n</pattern>
+ *     </encoder>
+ *   </appender>
+ *   <root level="DEBUG">
+ *     <appender-ref ref="TEST" />
+ *     <appender-ref ref="STDOUT" />
+ *   </root>
+ * </configuration>
+ * }</pre>
+ */
+public final class TestLogbackAppender extends AppenderBase<ILoggingEvent> {
+
+  @Override
+  protected void append(final ILoggingEvent event) {
+    final CapturedLog log =
+        new CapturedLog(
+            event.getMarker(),
+            event.getLevel().levelStr,
+            event.getMessage(),
+            event.getArgumentArray(),
+            event.getFormattedMessage());
+    TestLogCollector.INSTANCE.addLog(log);
+  }
+}