Improve Config JUnit extension

utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfig.java

@@ -1,9 +1,11 @@
 package datadog.trace.junit.utils.config;
 
-import java.lang.annotation.ElementType;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.TYPE;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
 import java.lang.annotation.Repeatable;
 import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 import org.junit.jupiter.api.extension.ExtendWith;
 
@@ -27,8 +29,8 @@
  * }
  * }</pre>
  */
-@Retention(RetentionPolicy.RUNTIME)
-@Target({ElementType.TYPE, ElementType.METHOD})
+@Retention(RUNTIME)
+@Target({TYPE, METHOD})
 @Repeatable(WithConfigs.class)
 @ExtendWith(WithConfigExtension.class)
 public @interface WithConfig {

utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigExtension.java

@@ -19,6 +19,7 @@
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -62,6 +63,7 @@ public class WithConfigExtension
   private static Field configInstanceField;
   private static Constructor<?> configConstructor;
 
+  private static volatile boolean configTransformerInstalled = false;
   private static volatile boolean isConfigInstanceModifiable = false;
   private static volatile boolean configModificationFailed = false;
 
@@ -73,22 +75,42 @@ public class WithConfigExtension
 
   @Override
   public void beforeAll(ExtensionContext context) {
-    installConfigTransformer();
+    /*
+     * Patch config classes to make them modifiable.
+     */
+    // Install config transformer error listener
+    if (!configTransformerInstalled) {
+      installConfigTransformer();
+      configTransformerInstalled = true;
+    }
+    // Make config instance modifiable
     makeConfigInstanceModifiable();
+    // Verify that config class transformation succeeded
     assertFalse(configModificationFailed, "Config class modification failed");
+    if (isConfigInstanceModifiable) {
+      checkConfigTransformation();
+    }
+    /*
+     * Back up config and apply class-level config values.
+     */
     if (originalSystemProperties == null) {
       saveProperties();
     }
+    // Apply class-level @WithConfig so config is available before @BeforeAll methods
+    applyClassLevelConfig(context);
+    if (isConfigInstanceModifiable) {
+      rebuildConfig();
+    }
   }
 
   @Override
   public void beforeEach(ExtensionContext context) {
     restoreProperties();
     environmentVariables.clear();
+    applyDeclaredConfig(context);
     if (isConfigInstanceModifiable) {
       rebuildConfig();
     }
-    applyDeclaredConfig(context);
   }
 
   @Override
@@ -108,14 +130,29 @@ public void afterAll(ExtensionContext context) {
     }
   }
 
-  private void applyDeclaredConfig(ExtensionContext context) {
-    // Class-level @WithConfig annotations (supports composed/meta-annotations)
-    List<WithConfig> classConfigs =
-        AnnotationSupport.findRepeatableAnnotations(
-            context.getRequiredTestClass(), WithConfig.class);
-    for (WithConfig cfg : classConfigs) {
-      applyConfig(cfg);
+  private static void applyDeclaredConfig(ExtensionContext context) {
+    applyClassLevelConfig(context);
+    applyMethodLevelConfig(context);
+  }
+
+  private static void applyClassLevelConfig(ExtensionContext context) {
+    // Walk the entire class hierarchy so annotations on superclasses and apply topmost first, then
+    // subclass overrides.
+    Class<?> testClass = context.getRequiredTestClass();
+    List<Class<?>> hierarchy = new ArrayList<>();
+    for (Class<?> cls = testClass; cls != null; cls = cls.getSuperclass()) {
+      hierarchy.add(cls);
+    }
+    for (int i = hierarchy.size() - 1; i >= 0; i--) {
+      List<WithConfig> classConfigs =
+          AnnotationSupport.findRepeatableAnnotations(hierarchy.get(i), WithConfig.class);
+      for (WithConfig cfg : classConfigs) {
+        applyConfig(cfg);
+      }
     }
+  }
+
+  private static void applyMethodLevelConfig(ExtensionContext context) {
     // Method-level @WithConfig annotations (supports composed/meta-annotations)
     context
         .getTestMethod()
@@ -131,12 +168,22 @@ private void applyDeclaredConfig(ExtensionContext context) {
 
   private static void applyConfig(WithConfig cfg) {
     if (cfg.env()) {
-      injectEnvConfig(cfg.key(), cfg.value(), cfg.addPrefix());
+      setEnvVariable(cfg.key(), cfg.value(), cfg.addPrefix());
     } else {
-      injectSysConfig(cfg.key(), cfg.value(), cfg.addPrefix());
+      setSysProperty(cfg.key(), cfg.value(), cfg.addPrefix());
     }
   }
 
+  private static void setSysProperty(String name, String value, boolean addPrefix) {
+    String prefixedName = addPrefix && !name.startsWith("dd.") ? "dd." + name : name;
+    System.setProperty(prefixedName, value);
+  }
+
+  private static void setEnvVariable(String name, String value, boolean addPrefix) {
+    String prefixedName = addPrefix && !name.startsWith("DD_") ? "DD_" + name : name;
+    environmentVariables.set(prefixedName, value);
+  }
+
   // endregion
 
   // region Public static API for imperative config injection
@@ -146,9 +193,7 @@ public static void injectSysConfig(String name, String value) {
   }
 
   public static void injectSysConfig(String name, String value, boolean addPrefix) {
-    checkConfigTransformation();
-    String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name;
-    System.setProperty(prefixedName, value);
+    setSysProperty(name, value, addPrefix);
     rebuildConfig();
   }
 
@@ -157,8 +202,7 @@ public static void removeSysConfig(String name) {
   }
 
   public static void removeSysConfig(String name, boolean addPrefix) {
-    checkConfigTransformation();
-    String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name;
+    String prefixedName = addPrefix && !name.startsWith("dd.") ? "dd." + name : name;
     System.clearProperty(prefixedName);
     rebuildConfig();
   }
@@ -168,9 +212,7 @@ public static void injectEnvConfig(String name, String value) {
   }
 
   public static void injectEnvConfig(String name, String value, boolean addPrefix) {
-    checkConfigTransformation();
-    String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name;
-    environmentVariables.set(prefixedName, value);
+    setEnvVariable(name, value, addPrefix);
     rebuildConfig();
   }
 
@@ -179,8 +221,7 @@ public static void removeEnvConfig(String name) {
   }
 
   public static void removeEnvConfig(String name, boolean addPrefix) {
-    checkConfigTransformation();
-    String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name;
+    String prefixedName = addPrefix && !name.startsWith("DD_") ? "DD_" + name : name;
     environmentVariables.removePrefixed(prefixedName);
     rebuildConfig();
   }
@@ -245,7 +286,6 @@ static void makeConfigInstanceModifiable() {
 
   private static void rebuildConfig() {
     synchronized (WithConfigExtension.class) {
-      checkConfigTransformation();
       try {
         Object newInstConfig = instConfigConstructor.newInstance();
         instConfigInstanceField.set(null, newInstConfig);

utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigs.java

@@ -1,14 +1,16 @@
 package datadog.trace.junit.utils.config;
 
-import java.lang.annotation.ElementType;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.TYPE;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
 import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 import org.junit.jupiter.api.extension.ExtendWith;
 
 /** Container annotation for repeatable {@link WithConfig}. */
-@Retention(RetentionPolicy.RUNTIME)
-@Target({ElementType.TYPE, ElementType.METHOD})
+@Retention(RUNTIME)
+@Target({TYPE, METHOD})
 @ExtendWith(WithConfigExtension.class)
 public @interface WithConfigs {
   WithConfig[] value();

utils/test-utils/src/main/java/datadog/trace/test/util/CleanConfigStateExtension.java

@@ -0,0 +1,80 @@
+package datadog.trace.test.util;
+
+import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
+
+import datadog.environment.EnvironmentVariables;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.TreeMap;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+/**
+ * Asserts that no {@code DD_*} environment variable and no {@code dd.*} system property (minus a
+ * small allowlist) is set around a test class.
+ */
+@SuppressForbidden
+public class CleanConfigStateExtension implements BeforeAllCallback, AfterAllCallback {
+
+  private static final List<String> ALLOWED_SYS_PROPS =
+      Arrays.asList(
+          "dd.appsec.enabled", "dd.iast.enabled", "dd.integration.grizzly-filterchain.enabled");
+
+  private static final Predicate<String> DATADOG_ENV_VAR_FILTER = k -> k.startsWith("DD_");
+  private static final Predicate<Object> DATADOG_SYS_PROPERTIES_FILTER =
+      o -> {
+        String key = (String) o;
+        return key.startsWith("DD_") && !ALLOWED_SYS_PROPS.contains(key);
+      };
+
+  @Override
+  public void beforeAll(ExtensionContext context) {
+    assertClean("before");
+  }
+
+  @Override
+  public void afterAll(ExtensionContext context) {
+    assertClean("after");
+  }
+
+  private static void assertClean(String phase) {
+    Map<String, String> leakedEnv =
+        filterMap(EnvironmentVariables.getAll(), DATADOG_ENV_VAR_FILTER);
+    Map<Object, Object> leakedSys =
+        filterMap(System.getProperties(), DATADOG_SYS_PROPERTIES_FILTER);
+    if (!leakedEnv.isEmpty() || !leakedSys.isEmpty()) {
+      assertionFailure()
+          .message("Leaked Datadog configuration detected " + phase + " test class")
+          .reason(formatLeaks(leakedEnv, leakedSys))
+          .buildAndThrow();
+    }
+  }
+
+  private static <T> Map<T, T> filterMap(Map<T, T> map, Predicate<T> keyFilter) {
+    return map.entrySet().stream()
+        .filter(e -> keyFilter.test(e.getKey()))
+        .collect(Collectors.toMap(Entry::getKey, Entry::getValue, (a, b) -> a, TreeMap::new));
+  }
+
+  private static String formatLeaks(Map<String, String> env, Map<Object, Object> sys) {
+    StringBuilder sb = new StringBuilder();
+    if (!env.isEmpty()) {
+      sb.append("environment variables:");
+      env.forEach((k, v) -> sb.append("\n  ").append(k).append('=').append(v));
+    }
+    if (!sys.isEmpty()) {
+      if (sb.length() > 0) {
+        sb.append('\n');
+      }
+      sb.append("system properties:");
+      sys.forEach((k, v) -> sb.append("\n  ").append(k).append('=').append(v));
+    }
+    return sb.toString();
+  }
+}

utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java

@@ -1,22 +1,21 @@
 package datadog.trace.test.util;
 
-import static org.junit.jupiter.api.Assertions.assertTrue;
-
-import datadog.environment.EnvironmentVariables;
 import datadog.trace.junit.utils.config.WithConfigExtension;
 import datadog.trace.junit.utils.context.AllowContextTestingExtension;
 import de.thetaphi.forbiddenapis.SuppressForbidden;
-import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.extension.ExtendWith;
 
-@ExtendWith({WithConfigExtension.class, AllowContextTestingExtension.class})
+@ExtendWith({
+  CleanConfigStateExtension.class,
+  WithConfigExtension.class,
+  AllowContextTestingExtension.class
+})
 @SuppressForbidden
 public class DDJavaSpecification {
 
@@ -27,13 +26,6 @@ public class DDJavaSpecification {
 
   @BeforeAll
   static void beforeAll() {
-    assertTrue(
-        EnvironmentVariables.getAll().entrySet().stream()
-            .noneMatch(e -> e.getKey().startsWith("DD_")));
-    assertTrue(
-        systemPropertiesExceptAllowed().entrySet().stream()
-            .noneMatch(e -> e.getKey().toString().startsWith("dd.")));
-
     if (getDDThreads().isEmpty()) {
       ignoreThreadCleanup = false;
     } else {
@@ -45,25 +37,9 @@ static void beforeAll() {
 
   @AfterAll
   static void afterAll() {
-    assertTrue(
-        EnvironmentVariables.getAll().entrySet().stream()
-            .noneMatch(e -> e.getKey().startsWith("DD_")));
-    assertTrue(
-        systemPropertiesExceptAllowed().entrySet().stream()
-            .noneMatch(e -> e.getKey().toString().startsWith("dd.")));
-
     checkThreads();
   }
 
-  private static Map<Object, Object> systemPropertiesExceptAllowed() {
-    List<String> allowlist =
-        Arrays.asList(
-            "dd.appsec.enabled", "dd.iast.enabled", "dd.integration.grizzly-filterchain.enabled");
-    return System.getProperties().entrySet().stream()
-        .filter(e -> !allowlist.contains(String.valueOf(e.getKey())))
-        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
-  }
-
   @AfterEach
   void cleanup() {
     if (assertThreadsEachCleanup) {