Add functional class injection tests using startup instrumentation

buildSrc/src/main/kotlin/instrument-testing.gradle.kts

@@ -0,0 +1,22 @@
+plugins {
+  id("java-common")
+}
+
+val testAgent = configurations.create("testAgent") {
+  isCanBeConsumed = false
+  isTransitive = false
+}
+
+dependencies {
+  testImplementation(project(":testing"))
+  testAgent(project(":testing"))
+}
+
+class TestAgentArguments(@get:Classpath val agentClasspath: FileCollection) : CommandLineArgumentProvider {
+  override fun asArguments() = listOf("-javaagent:${agentClasspath.singleFile.absolutePath}")
+}
+
+// Supply instrumentation to every test JVM, including the additional Java version tasks.
+tasks.withType<Test>().configureEach {
+  jvmArgumentProviders.add(TestAgentArguments(testAgent))
+}

class-inject/build.gradle.kts

@@ -1,6 +1,7 @@
 plugins {
   id("java-common")
   id("instrument-glue")
+  id("instrument-testing")
 }
 
 dependencies {

class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java

@@ -0,0 +1,139 @@
+package datadog.instrument.classinject;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.objectweb.asm.Opcodes.*;
+
+import datadog.instrument.testing.TestAgent;
+import java.security.ProtectionDomain;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.MethodVisitor;
+
+class ClassInjectorTest {
+
+  @BeforeAll
+  static void enableClassInjection() {
+    ClassInjector.enableClassInjection(TestAgent.getInstrumentation());
+  }
+
+  @Test
+  void injectsBootClass() throws Exception {
+    String name = "datadog.injected.BootSample";
+    Class<?> injectedClass = ClassInjector.injectBootClass(name, sampleClass(name));
+    assertThat(injectedClass.getClassLoader()).isNull();
+    assertThat(Class.forName(name, false, null)).isSameAs(injectedClass);
+    assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42);
+
+    // injecting an already defined class returns the existing class
+    assertThat(ClassInjector.injectBootClass(name, sampleClass(name))).isSameAs(injectedClass);
+  }
+
+  @Test
+  void injectsClassIntoGivenLoader() throws Exception {
+    ClassLoader loader = newCL();
+    String name = "datadog.injected.LoaderSample";
+    Class<?> injectedClass = ClassInjector.injectClass(name, sampleClass(name), loader);
+    assertThat(injectedClass.getClassLoader()).isSameAs(loader);
+    assertThat(loader.loadClass(name)).isSameAs(injectedClass);
+    assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42);
+
+    // injecting an already defined class returns the existing class
+    assertThat(ClassInjector.injectClass(name, sampleClass(name), loader)).isSameAs(injectedClass);
+  }
+
+  @Test
+  void injectsClassWithGivenProtectionDomain() throws Exception {
+    ClassLoader loader = newCL();
+    ProtectionDomain domain = new ProtectionDomain(null, null, loader, null);
+    String name = "datadog.injected.DomainSample";
+    Class<?> injectedClass = ClassInjector.injectClass(name, sampleClass(name), domain);
+    assertThat(injectedClass.getClassLoader()).isSameAs(loader);
+    assertThat(injectedClass.getProtectionDomain()).isSameAs(domain);
+    assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42);
+
+    assertThat(ClassInjector.injectClass(name, sampleClass(name), domain)).isSameAs(injectedClass);
+  }
+
+  @Test
+  void injectsBootClassWithGivenProtectionDomain() throws Exception {
+    ProtectionDomain domain = new ProtectionDomain(null, null, null, null);
+    String name = "datadog.injected.BootDomainSample";
+    Class<?> injectedClass = ClassInjector.injectClass(name, sampleClass(name), domain);
+    assertThat(injectedClass.getClassLoader()).isNull();
+    assertThat(Class.forName(name, false, null)).isSameAs(injectedClass);
+    assertThat(injectedClass.getProtectionDomain()).isSameAs(domain);
+    assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42);
+
+    assertThat(ClassInjector.injectClass(name, sampleClass(name), domain)).isSameAs(injectedClass);
+  }
+
+  @Test
+  void injectsBootClassesThatReferenceEachOther() throws Exception {
+    String caller = "datadog.injected.BootCallerSample";
+    String callee = "datadog.injected.BootCalleeSample";
+    Map<String, byte[]> bytecode = new LinkedHashMap<>();
+    bytecode.put(caller, delegatingClass(caller, callee));
+    bytecode.put(callee, sampleClass(callee));
+
+    List<Class<?>> injectedClasses = ClassInjector.injectBootClasses(bytecode);
+    assertThat(injectedClasses).extracting(Class::getName).containsExactly(caller, callee);
+    assertThat(injectedClasses).allSatisfy(c -> assertThat(c.getClassLoader()).isNull());
+    assertThat(Class.forName(caller, false, null)).isSameAs(injectedClasses.get(0));
+    assertThat(Class.forName(callee, false, null)).isSameAs(injectedClasses.get(1));
+    assertThat(injectedClasses.get(0).getMethod("value").invoke(null)).isEqualTo(42);
+
+    assertThat(ClassInjector.injectBootClasses(bytecode))
+        .containsExactlyElementsOf(injectedClasses);
+  }
+
+  @Test
+  void injectsClassesThatReferenceEachOther() throws Exception {
+    ClassLoader loader = newCL();
+    String caller = "datadog.injected.CallerSample";
+    String callee = "datadog.injected.CalleeSample";
+    Map<String, byte[]> bytecode = new LinkedHashMap<>();
+    bytecode.put(caller, delegatingClass(caller, callee));
+    bytecode.put(callee, sampleClass(callee));
+
+    List<Class<?>> injectedClasses = ClassInjector.injectClasses(bytecode, loader);
+    assertThat(injectedClasses).extracting(Class::getName).containsExactly(caller, callee);
+    assertThat(injectedClasses).allSatisfy(c -> assertThat(c.getClassLoader()).isSameAs(loader));
+    assertThat(injectedClasses.get(0).getMethod("value").invoke(null)).isEqualTo(42);
+  }
+
+  private static ClassLoader newCL() {
+    return new ClassLoader(null) {};
+  }
+
+  /** Generates a class with a static {@code value()} method that returns 42. */
+  private static byte[] sampleClass(String name) {
+    ClassWriter writer = new ClassWriter(0);
+    writer.visit(V1_8, ACC_PUBLIC, name.replace('.', '/'), null, "java/lang/Object", null);
+    MethodVisitor method = writer.visitMethod(ACC_PUBLIC | ACC_STATIC, "value", "()I", null, null);
+    method.visitCode();
+    method.visitIntInsn(BIPUSH, 42);
+    method.visitInsn(IRETURN);
+    method.visitMaxs(1, 0);
+    method.visitEnd();
+    writer.visitEnd();
+    return writer.toByteArray();
+  }
+
+  /** Generates a class with a static {@code value()} method that delegates to the target. */
+  private static byte[] delegatingClass(String name, String target) {
+    ClassWriter writer = new ClassWriter(0);
+    writer.visit(V1_8, ACC_PUBLIC, name.replace('.', '/'), null, "java/lang/Object", null);
+    MethodVisitor method = writer.visitMethod(ACC_PUBLIC | ACC_STATIC, "value", "()I", null, null);
+    method.visitCode();
+    method.visitMethodInsn(INVOKESTATIC, target.replace('.', '/'), "value", "()I", false);
+    method.visitInsn(IRETURN);
+    method.visitMaxs(1, 0);
+    method.visitEnd();
+    writer.visitEnd();
+    return writer.toByteArray();
+  }
+}

gradle.properties

@@ -4,6 +4,6 @@ org.gradle.parallel=true
 org.gradle.workers.max=4
 org.gradle.java.installations.auto-detect=false
 org.gradle.java.installations.auto-download=false
-org.gradle.java.installations.fromEnv=JAVA_8_HOME,JAVA_11_HOME,JAVA_17_HOME,JAVA_21_HOME,JAVA_25_HOME
+org.gradle.java.installations.fromEnv=JAVA_8_HOME,JAVA_11_HOME,JAVA_17_HOME,JAVA_21_HOME,JAVA_25_HOME,JAVA_27_HOME
 org.gradle.problems.report=false
 org.gradle.warning.mode=none

testing/build.gradle.kts

@@ -1,3 +1,12 @@
 plugins {
   id("java-common")
 }
+
+// Allow test JVMs to obtain instrumentation without dynamic agent attachment.
+tasks.jar {
+  manifest {
+    attributes(
+      "Premain-Class" to "datadog.instrument.testing.TestAgent",
+      "Can-Retransform-Classes" to "true")
+  }
+}

testing/src/main/java/datadog/instrument/testing/TestAgent.java

@@ -0,0 +1,29 @@
+/*
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/).
+ * Copyright 2026-Present Datadog, Inc.
+ */
+
+package datadog.instrument.testing;
+
+import java.lang.instrument.Instrumentation;
+
+/** Supplies instrumentation to tests when the testing jar is used as a startup agent. */
+public final class TestAgent {
+  private static volatile Instrumentation instrumentation;
+
+  private TestAgent() {}
+
+  /** Called by the JVM when the testing jar is supplied with {@code -javaagent}. */
+  public static void premain(String arguments, Instrumentation inst) {
+    instrumentation = inst;
+  }
+
+  /** Returns the instrumentation supplied at JVM startup. */
+  public static Instrumentation getInstrumentation() {
+    if (instrumentation == null) {
+      throw new IllegalStateException("The testing jar must be supplied with -javaagent");
+    }
+    return instrumentation;
+  }
+}