From f588bcf249b0d540524595216f20427c7a01bc45 Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Fri, 10 Apr 2026 09:41:02 +0200 Subject: [PATCH 1/4] feat(junit): Add config extension --- settings.gradle.kts | 1 + utils/junit-utils/build.gradle.kts | 14 + .../trace/junit/utils/config/WithConfig.java | 49 +++ .../utils/config/WithConfigExtension.java | 361 ++++++++++++++++++ .../trace/junit/utils/config/WithConfigs.java | 15 + utils/test-utils/build.gradle.kts | 1 + .../trace/test/util/DDJavaSpecification.java | 286 +------------- 7 files changed, 444 insertions(+), 283 deletions(-) create mode 100644 utils/junit-utils/build.gradle.kts create mode 100644 utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfig.java create mode 100644 utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigExtension.java create mode 100644 utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigs.java diff --git a/settings.gradle.kts b/settings.gradle.kts index 0bac052092a..086913ae2da 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -159,6 +159,7 @@ include( ":dd-java-agent:testing", ":utils:config-utils", ":utils:container-utils", + ":utils:junit-utils", ":utils:filesystem-utils", ":utils:flare-utils", ":utils:logging-utils", diff --git a/utils/junit-utils/build.gradle.kts b/utils/junit-utils/build.gradle.kts new file mode 100644 index 00000000000..e26ddcfa1ea --- /dev/null +++ b/utils/junit-utils/build.gradle.kts @@ -0,0 +1,14 @@ +plugins { + `java-library` +} + +apply(from = "$rootDir/gradle/java.gradle") + +dependencies { + api(libs.bytebuddy) + api(libs.bytebuddyagent) + api(libs.forbiddenapis) + api(project(":components:environment")) + + compileOnly(libs.junit.jupiter) +} diff --git a/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfig.java b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfig.java new file mode 100644 index 00000000000..172c24f03ff --- /dev/null +++ b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfig.java @@ -0,0 +1,49 @@ +package datadog.trace.junit.utils.config; + +import java.lang.annotation.ElementType; +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; + +/** + * Declares a configuration override for a test. Can be placed on a test class (applies to all + * tests) or on individual test methods. + * + *

By default, injects a system property with the {@code dd.} prefix. Use {@code env = true} for + * environment variables (prefix {@code DD_}). + * + *

Examples: + * + * 

{@code
+ * @WithConfig(key = "service", value = "my_service")
+ * @WithConfig(key = "trace.resolver.enabled", value = "false")
+ * class MyTest extends DDJavaSpecification {
+ *
+ *   @Test
+ *   @WithConfig(key = "AGENT_HOST", value = "localhost", env = true)
+ *   void testWithEnv() { ... }
+ * }
+ * }
+ */ +@Retention(RetentionPolicy.RUNTIME) +@Target({ElementType.TYPE, ElementType.METHOD}) +@Repeatable(WithConfigs.class) +@ExtendWith(WithConfigExtension.class) +public @interface WithConfig { + /** + * Config key (e.g. {@code "trace.resolver.enabled"}). The {@code dd.}/{@code DD_} prefix is + * auto-added unless {@link #addPrefix()} is {@code false}. + */ + String key(); + + /** Config value. */ + String value(); + + /** If {@code true}, sets an environment variable instead of a system property. */ + boolean env() default false; + + /** If {@code false}, the key is used as-is without adding the {@code dd.}/{@code DD_} prefix. */ + boolean addPrefix() default true; +} diff --git a/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigExtension.java b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigExtension.java new file mode 100644 index 00000000000..5041d8dc669 --- /dev/null +++ b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigExtension.java @@ -0,0 +1,361 @@ +package datadog.trace.junit.utils.config; + +import static net.bytebuddy.agent.builder.AgentBuilder.RedefinitionStrategy.Listener.ErrorEscalating.FAIL_FAST; +import static net.bytebuddy.agent.builder.AgentBuilder.RedefinitionStrategy.RETRANSFORMATION; +import static net.bytebuddy.description.modifier.FieldManifestation.VOLATILE; +import static net.bytebuddy.description.modifier.Ownership.STATIC; +import static net.bytebuddy.description.modifier.Visibility.PUBLIC; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.namedOneOf; +import static net.bytebuddy.matcher.ElementMatchers.none; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.environment.EnvironmentVariables; +import de.thetaphi.forbiddenapis.SuppressForbidden; +import edu.umd.cs.findbugs.annotations.NonNull; +import java.lang.instrument.Instrumentation; +import java.lang.reflect.Constructor; +import java.lang.reflect.Field; +import java.lang.reflect.Modifier; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import net.bytebuddy.agent.ByteBuddyAgent; +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.dynamic.ClassFileLocator; +import net.bytebuddy.dynamic.Transformer; +import net.bytebuddy.utility.JavaModule; +import org.junit.jupiter.api.extension.AfterAllCallback; +import org.junit.jupiter.api.extension.AfterEachCallback; +import org.junit.jupiter.api.extension.BeforeAllCallback; +import org.junit.jupiter.api.extension.BeforeEachCallback; +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.platform.commons.support.AnnotationSupport; + +/** + * JUnit 5 extension that manages DD config injection for tests. Handles: + * + * + * + *

This extension is auto-registered when using {@link WithConfig} annotations. It can also be + * used explicitly via {@code @ExtendWith(WithConfigExtension.class)}. + */ +@SuppressForbidden +public class WithConfigExtension + implements BeforeAllCallback, BeforeEachCallback, AfterEachCallback, AfterAllCallback { + + static final String INST_CONFIG = "datadog.trace.api.InstrumenterConfig"; + static final String CONFIG = "datadog.trace.api.Config"; + + private static Field instConfigInstanceField; + private static Constructor instConfigConstructor; + private static Field configInstanceField; + private static Constructor configConstructor; + + private static volatile boolean isConfigInstanceModifiable = false; + private static volatile boolean configModificationFailed = false; + + static final TestEnvironmentVariables environmentVariables = TestEnvironmentVariables.setup(); + + private static Properties originalSystemProperties; + + // region JUnit lifecycle callbacks + + @Override + public void beforeAll(ExtensionContext context) { + installConfigTransformer(); + makeConfigInstanceModifiable(); + assertFalse(configModificationFailed, "Config class modification failed"); + if (originalSystemProperties == null) { + saveProperties(); + } + } + + @Override + public void beforeEach(ExtensionContext context) { + restoreProperties(); + environmentVariables.clear(); + if (isConfigInstanceModifiable) { + rebuildConfig(); + } + applyDeclaredConfig(context); + } + + @Override + public void afterEach(ExtensionContext context) { + environmentVariables.clear(); + restoreProperties(); + if (isConfigInstanceModifiable) { + rebuildConfig(); + } + } + + @Override + public void afterAll(ExtensionContext context) { + restoreProperties(); + if (isConfigInstanceModifiable) { + rebuildConfig(); + } + } + + private void applyDeclaredConfig(ExtensionContext context) { + // Class-level @WithConfig annotations (supports composed/meta-annotations) + List classConfigs = + AnnotationSupport.findRepeatableAnnotations( + context.getRequiredTestClass(), WithConfig.class); + for (WithConfig cfg : classConfigs) { + applyConfig(cfg); + } + // Method-level @WithConfig annotations (supports composed/meta-annotations) + context + .getTestMethod() + .ifPresent( + method -> { + List methodConfigs = + AnnotationSupport.findRepeatableAnnotations(method, WithConfig.class); + for (WithConfig cfg : methodConfigs) { + applyConfig(cfg); + } + }); + } + + private static void applyConfig(WithConfig cfg) { + if (cfg.env()) { + injectEnvConfig(cfg.key(), cfg.value(), cfg.addPrefix()); + } else { + injectSysConfig(cfg.key(), cfg.value(), cfg.addPrefix()); + } + } + + // endregion + + // region Public static API for imperative config injection + + public static void injectSysConfig(String name, String value) { + injectSysConfig(name, value, true); + } + + public static void injectSysConfig(String name, String value, boolean addPrefix) { + checkConfigTransformation(); + String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name; + System.setProperty(prefixedName, value); + rebuildConfig(); + } + + public static void removeSysConfig(String name) { + removeSysConfig(name, true); + } + + public static void removeSysConfig(String name, boolean addPrefix) { + checkConfigTransformation(); + String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name; + System.clearProperty(prefixedName); + rebuildConfig(); + } + + public static void injectEnvConfig(String name, String value) { + injectEnvConfig(name, value, true); + } + + public static void injectEnvConfig(String name, String value, boolean addPrefix) { + checkConfigTransformation(); + String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name; + environmentVariables.set(prefixedName, value); + rebuildConfig(); + } + + public static void removeEnvConfig(String name) { + removeEnvConfig(name, true); + } + + public static void removeEnvConfig(String name, boolean addPrefix) { + checkConfigTransformation(); + String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name; + environmentVariables.removePrefixed(prefixedName); + rebuildConfig(); + } + + // endregion + + // region Config infrastructure setup + + private static void installConfigTransformer() { + try { + Instrumentation instrumentation = ByteBuddyAgent.install(); + new AgentBuilder.Default() + .with(RETRANSFORMATION) + .with(FAIL_FAST) + .with( + new AgentBuilder.LocationStrategy.Simple( + ClassFileLocator.ForClassLoader.ofSystemLoader())) + .ignore(none()) + .type(namedOneOf(INST_CONFIG, CONFIG)) + .transform( + (builder, typeDescription, classLoader, module, pd) -> + builder + .field(named("INSTANCE")) + .transform(Transformer.ForField.withModifiers(PUBLIC, STATIC, VOLATILE))) + .with(new ConfigInstrumentationFailedListener()) + .installOn(instrumentation); + } catch (IllegalStateException e) { + // Ignore. When we have -javaagent:dd-java-agent.jar, this is fine. + } + } + + static void makeConfigInstanceModifiable() { + if (isConfigInstanceModifiable || configModificationFailed) { + return; + } + + try { + Class instConfigClass = Class.forName(INST_CONFIG); + instConfigInstanceField = instConfigClass.getDeclaredField("INSTANCE"); + instConfigConstructor = instConfigClass.getDeclaredConstructor(); + instConfigConstructor.setAccessible(true); + Class configClass = Class.forName(CONFIG); + configInstanceField = configClass.getDeclaredField("INSTANCE"); + configConstructor = configClass.getDeclaredConstructor(); + configConstructor.setAccessible(true); + + isConfigInstanceModifiable = true; + } catch (ClassNotFoundException e) { + if (INST_CONFIG.equals(e.getMessage()) || CONFIG.equals(e.getMessage())) { + System.err.println("Config class not found in this classloader. Not transforming it"); + } else { + configModificationFailed = true; + System.err.println("Config will not be modifiable"); + e.printStackTrace(); + } + } catch (ReflectiveOperationException e) { + configModificationFailed = true; + System.err.println("Config will not be modifiable"); + e.printStackTrace(); + } + } + + private static void rebuildConfig() { + synchronized (WithConfigExtension.class) { + checkConfigTransformation(); + try { + Object newInstConfig = instConfigConstructor.newInstance(); + instConfigInstanceField.set(null, newInstConfig); + Object newConfig = configConstructor.newInstance(); + configInstanceField.set(null, newConfig); + } catch (ReflectiveOperationException e) { + throw new AssertionError("Failed to rebuild config", e); + } + } + } + + // endregion + + // region Property management + + static void saveProperties() { + originalSystemProperties = new Properties(); + originalSystemProperties.putAll(System.getProperties()); + } + + static void restoreProperties() { + if (originalSystemProperties != null) { + Properties copy = new Properties(); + copy.putAll(originalSystemProperties); + System.setProperties(copy); + } + } + + // endregion + + // region Validation + + private static void checkConfigTransformation() { + assertTrue(isConfigInstanceModifiable); + assertNotNull(instConfigConstructor); + checkWritable(instConfigInstanceField); + assertNotNull(configConstructor); + checkWritable(configInstanceField); + } + + private static void checkWritable(Field field) { + assertNotNull(field); + assertTrue(Modifier.isPublic(field.getModifiers())); + assertTrue(Modifier.isStatic(field.getModifiers())); + assertTrue(Modifier.isVolatile(field.getModifiers())); + assertFalse(Modifier.isFinal(field.getModifiers())); + } + + // endregion + + /** Test-only environment variable provider that replaces the real one during tests. */ + public static class TestEnvironmentVariables + extends EnvironmentVariables.EnvironmentVariablesProvider { + private final Map env = new HashMap<>(); + + TestEnvironmentVariables(String... kv) { + for (int i = 0; i + 1 < kv.length; i += 2) { + this.env.put(kv[i], kv[i + 1]); + } + } + + @Override + public String get(@NonNull String name) { + return env.get(name); + } + + @Override + public Map getAll() { + return env; + } + + public void set(String name, String value) { + env.put(name, value); + } + + public void removePrefixed(String prefix) { + env.keySet().removeIf(k -> k.startsWith(prefix)); + } + + public void clear() { + env.clear(); + } + + @SuppressForbidden + static TestEnvironmentVariables setup(String... kv) { + TestEnvironmentVariables provider = new TestEnvironmentVariables(kv); + EnvironmentVariables.provider = provider; + + String propagateVars = System.getenv("TEST_ENV_PROPAGATE_VARS"); + if (propagateVars != null) { + for (String envVar : propagateVars.split(",")) { + provider.env.put(envVar, System.getenv(envVar)); + } + } + + return provider; + } + } + + private static class ConfigInstrumentationFailedListener extends AgentBuilder.Listener.Adapter { + @Override + public void onError( + @NonNull String typeName, + ClassLoader classLoader, + JavaModule module, + boolean loaded, + @NonNull Throwable throwable) { + if (CONFIG.equals(typeName)) { + configModificationFailed = true; + } + } + } +} diff --git a/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigs.java b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigs.java new file mode 100644 index 00000000000..64a21fe5845 --- /dev/null +++ b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfigs.java @@ -0,0 +1,15 @@ +package datadog.trace.junit.utils.config; + +import java.lang.annotation.ElementType; +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}) +@ExtendWith(WithConfigExtension.class) +public @interface WithConfigs { + WithConfig[] value(); +} diff --git a/utils/test-utils/build.gradle.kts b/utils/test-utils/build.gradle.kts index 48b3dc39403..07b6ccf1e3f 100644 --- a/utils/test-utils/build.gradle.kts +++ b/utils/test-utils/build.gradle.kts @@ -11,6 +11,7 @@ dependencies { api(project(":components:environment")) api(project(":utils:config-utils")) + api(project(":utils:junit-utils")) api(group = "commons-fileupload", name = "commons-fileupload", version = "1.5") compileOnly(libs.junit.jupiter) diff --git a/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java b/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java index 8119ad879e4..028ec3a7bce 100644 --- a/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java +++ b/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java @@ -1,39 +1,22 @@ package datadog.trace.test.util; -import static net.bytebuddy.description.modifier.FieldManifestation.VOLATILE; -import static net.bytebuddy.description.modifier.Ownership.STATIC; -import static net.bytebuddy.description.modifier.Visibility.PUBLIC; -import static net.bytebuddy.matcher.ElementMatchers.named; -import static net.bytebuddy.matcher.ElementMatchers.namedOneOf; -import static net.bytebuddy.matcher.ElementMatchers.none; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import datadog.environment.EnvironmentVariables; +import datadog.trace.junit.utils.config.WithConfigExtension; import de.thetaphi.forbiddenapis.SuppressForbidden; -import java.lang.instrument.Instrumentation; -import java.lang.reflect.Constructor; -import java.lang.reflect.Field; import java.lang.reflect.Method; -import java.lang.reflect.Modifier; import java.util.Arrays; -import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Properties; import java.util.Set; import java.util.stream.Collectors; -import net.bytebuddy.agent.ByteBuddyAgent; -import net.bytebuddy.agent.builder.AgentBuilder; -import net.bytebuddy.dynamic.ClassFileLocator; -import net.bytebuddy.dynamic.Transformer; -import net.bytebuddy.utility.JavaModule; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.extension.ExtendWith; +@ExtendWith(WithConfigExtension.class) @SuppressForbidden public class DDJavaSpecification { @@ -41,22 +24,8 @@ public class DDJavaSpecification { static final String CONTEXT_BINDER = "datadog.context.ContextBinder"; static final String CONTEXT_MANAGER = "datadog.context.ContextManager"; - static final String INST_CONFIG = "datadog.trace.api.InstrumenterConfig"; - static final String CONFIG = "datadog.trace.api.Config"; - - private static Field instConfigInstanceField; - private static Constructor instConfigConstructor; - private static Field configInstanceField; - private static Constructor configConstructor; private static Boolean contextTestingAllowed; - private static volatile boolean isConfigInstanceModifiable = false; - static volatile boolean configModificationFailed = false; - - protected static final TestEnvironmentVariables environmentVariables = - TestEnvironmentVariables.setup(); - - private static Properties originalSystemProperties; protected boolean assertThreadsEachCleanup = true; private static volatile boolean ignoreThreadCleanup; @@ -64,11 +33,6 @@ public class DDJavaSpecification { @BeforeAll static void beforeAll() { allowContextTesting(); - installConfigTransformer(); - makeConfigInstanceModifiable(); - assertFalse( - configModificationFailed, - "Config class modification failed. Ensure all test classes extend DDJavaSpecification"); assertTrue( EnvironmentVariables.getAll().entrySet().stream() .noneMatch(e -> e.getKey().startsWith("DD_"))); @@ -86,7 +50,6 @@ static void beforeAll() { "Found DD threads before test started. Ignoring thread cleanup for this test class"); ignoreThreadCleanup = true; } - saveProperties(); } static void allowContextTesting() { @@ -110,77 +73,8 @@ static void allowContextTesting() { } } - private static void installConfigTransformer() { - try { - Instrumentation instrumentation = ByteBuddyAgent.install(); - new AgentBuilder.Default() - .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) - .with(AgentBuilder.RedefinitionStrategy.Listener.ErrorEscalating.FAIL_FAST) - .with( - new AgentBuilder.LocationStrategy.Simple( - ClassFileLocator.ForClassLoader.ofSystemLoader())) - .ignore(none()) - .type(namedOneOf(INST_CONFIG, CONFIG)) - .transform( - (builder, typeDescription, classLoader, module, pd) -> - builder - .field(named("INSTANCE")) - .transform(Transformer.ForField.withModifiers(PUBLIC, STATIC, VOLATILE))) - .with(new ConfigInstrumentationFailedListener()) - .installOn(instrumentation); - } catch (IllegalStateException e) { - // Ignore. When we have -javaagent:dd-java-agent.jar, this is fine. - } - } - - static void makeConfigInstanceModifiable() { - if (isConfigInstanceModifiable || configModificationFailed) { - return; - } - - try { - Class instConfigClass = Class.forName(INST_CONFIG); - instConfigInstanceField = instConfigClass.getDeclaredField("INSTANCE"); - instConfigConstructor = instConfigClass.getDeclaredConstructor(); - instConfigConstructor.setAccessible(true); - Class configClass = Class.forName(CONFIG); - configInstanceField = configClass.getDeclaredField("INSTANCE"); - configConstructor = configClass.getDeclaredConstructor(); - configConstructor.setAccessible(true); - - isConfigInstanceModifiable = true; - } catch (ClassNotFoundException e) { - if (INST_CONFIG.equals(e.getMessage()) || CONFIG.equals(e.getMessage())) { - System.out.println("Config class not found in this classloader. Not transforming it"); - } else { - configModificationFailed = true; - System.out.println("Config will not be modifiable"); - e.printStackTrace(); - } - } catch (ReflectiveOperationException e) { - configModificationFailed = true; - System.out.println("Config will not be modifiable"); - e.printStackTrace(); - } - } - - private static void saveProperties() { - originalSystemProperties = new Properties(); - originalSystemProperties.putAll(System.getProperties()); - } - - private static void restoreProperties() { - if (originalSystemProperties != null) { - Properties copy = new Properties(); - copy.putAll(originalSystemProperties); - System.setProperties(copy); - } - } - @AfterAll static void afterAll() { - restoreProperties(); - assertTrue( EnvironmentVariables.getAll().entrySet().stream() .noneMatch(e -> e.getKey().startsWith("DD_"))); @@ -188,10 +82,6 @@ static void afterAll() { systemPropertiesExceptAllowed().entrySet().stream() .noneMatch(e -> e.getKey().toString().startsWith("dd."))); - if (isConfigInstanceModifiable) { - rebuildConfig(); - } - checkThreads(); } @@ -204,39 +94,8 @@ private static Map systemPropertiesExceptAllowed() { .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } - @BeforeEach - void setup() { - restoreProperties(); - - assertTrue( - EnvironmentVariables.getAll().entrySet().stream() - .noneMatch(e -> e.getKey().startsWith("DD_"))); - assertTrue( - systemPropertiesExceptAllowed().entrySet().stream() - .noneMatch(e -> e.getKey().toString().startsWith("dd."))); - - if (isConfigInstanceModifiable) { - rebuildConfig(); - } - } - @AfterEach void cleanup() { - environmentVariables.clear(); - - restoreProperties(); - - assertTrue( - EnvironmentVariables.getAll().entrySet().stream() - .noneMatch(e -> e.getKey().startsWith("DD_"))); - assertTrue( - systemPropertiesExceptAllowed().entrySet().stream() - .noneMatch(e -> e.getKey().toString().startsWith("dd."))); - - if (isConfigInstanceModifiable) { - rebuildConfig(); - } - if (assertThreadsEachCleanup) { checkThreads(); } @@ -277,143 +136,4 @@ static void checkThreads() { } } - public void injectSysConfig(String name, String value) { - injectSysConfig(name, value, true); - } - - public void injectSysConfig(String name, String value, boolean addPrefix) { - checkConfigTransformation(); - - String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name; - System.setProperty(prefixedName, value); - rebuildConfig(); - } - - public void removeSysConfig(String name) { - removeSysConfig(name, true); - } - - public void removeSysConfig(String name, boolean addPrefix) { - checkConfigTransformation(); - - String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name; - System.clearProperty(prefixedName); - rebuildConfig(); - } - - public void injectEnvConfig(String name, String value) { - injectEnvConfig(name, value, true); - } - - public void injectEnvConfig(String name, String value, boolean addPrefix) { - checkConfigTransformation(); - - String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name; - environmentVariables.set(prefixedName, value); - rebuildConfig(); - } - - public void removeEnvConfig(String name) { - removeEnvConfig(name, true); - } - - public void removeEnvConfig(String name, boolean addPrefix) { - checkConfigTransformation(); - - String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name; - environmentVariables.removePrefixed(prefixedName); - rebuildConfig(); - } - - static void rebuildConfig() { - synchronized (DDJavaSpecification.class) { - checkConfigTransformation(); - try { - Object newInstConfig = instConfigConstructor.newInstance(); - instConfigInstanceField.set(null, newInstConfig); - Object newConfig = configConstructor.newInstance(); - configInstanceField.set(null, newConfig); - } catch (ReflectiveOperationException e) { - throw new AssertionError("Failed to rebuild config", e); - } - } - } - - private static void checkConfigTransformation() { - assertTrue(isConfigInstanceModifiable); - assertNotNull(instConfigConstructor); - checkWritable(instConfigInstanceField); - assertNotNull(configConstructor); - checkWritable(configInstanceField); - } - - private static void checkWritable(Field field) { - assertNotNull(field); - assertTrue(Modifier.isPublic(field.getModifiers())); - assertTrue(Modifier.isStatic(field.getModifiers())); - assertTrue(Modifier.isVolatile(field.getModifiers())); - assertFalse(Modifier.isFinal(field.getModifiers())); - } - - public static class TestEnvironmentVariables - extends EnvironmentVariables.EnvironmentVariablesProvider { - private final Map env = new HashMap<>(); - - TestEnvironmentVariables(String... kv) { - for (int i = 0; i + 1 < kv.length; i += 2) { - env.put(kv[i], kv[i + 1]); - } - } - - @Override - public String get(String name) { - return env.get(name); - } - - @Override - public Map getAll() { - return env; - } - - public void set(String name, String value) { - env.put(name, value); - } - - public void removePrefixed(String prefix) { - env.keySet().removeIf(k -> k.startsWith(prefix)); - } - - public void clear() { - env.clear(); - } - - @SuppressForbidden - static TestEnvironmentVariables setup(String... kv) { - TestEnvironmentVariables provider = new TestEnvironmentVariables(kv); - EnvironmentVariables.provider = provider; - - String propagateVars = System.getenv("TEST_ENV_PROPAGATE_VARS"); - if (propagateVars != null) { - for (String envVar : propagateVars.split(",")) { - provider.env.put(envVar, System.getenv(envVar)); - } - } - - return provider; - } - } - - private static class ConfigInstrumentationFailedListener extends AgentBuilder.Listener.Adapter { - @Override - public void onError( - String typeName, - ClassLoader classLoader, - JavaModule module, - boolean loaded, - Throwable throwable) { - if (CONFIG.equals(typeName)) { - configModificationFailed = true; - } - } - } } From d243e8131fca8bad1631fd527ea13bb403398f3c Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Fri, 10 Apr 2026 10:24:22 +0200 Subject: [PATCH 2/4] feat(junit): Add context extension --- .../test/AbstractInstrumentationTest.java | 18 ++-------- .../utils/context/AllowContextTesting.java | 16 +++++++++ .../context/AllowContextTestingExtension.java | 27 ++++++++++++++ .../trace/test/util/DDJavaSpecification.java | 35 ++----------------- 4 files changed, 47 insertions(+), 49 deletions(-) create mode 100644 utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTesting.java create mode 100644 utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTestingExtension.java diff --git a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java index af7a7fa40d7..cfcc93ed006 100644 --- a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java +++ b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java @@ -20,7 +20,7 @@ import datadog.trace.core.DDSpan; import datadog.trace.core.PendingTrace; import datadog.trace.core.TraceCollector; -import de.thetaphi.forbiddenapis.SuppressForbidden; +import datadog.trace.junit.utils.context.AllowContextTestingExtension; import java.lang.instrument.ClassFileTransformer; import java.lang.instrument.Instrumentation; import java.util.List; @@ -31,7 +31,6 @@ import java.util.function.Predicate; import net.bytebuddy.agent.ByteBuddyAgent; import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.extension.ExtendWith; import org.opentest4j.AssertionFailedError; @@ -42,7 +41,7 @@ * current implementation is inspired and kept close to it Groovy / Spock counterpart, the {@code * InstrumentationSpecification}. */ -@ExtendWith(TestClassShadowingExtension.class) +@ExtendWith({TestClassShadowingExtension.class, AllowContextTestingExtension.class}) public abstract class AbstractInstrumentationTest { static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation(); @@ -55,19 +54,6 @@ public abstract class AbstractInstrumentationTest { protected ClassFileTransformer activeTransformer; protected ClassFileTransformerListener transformerLister; - @SuppressForbidden // Class.forName() used to dynamically configure context if present - @BeforeAll - static void allowContextTesting() { - // Allow re-registration of context managers so each test can use a fresh tracer. - // This mirrors DDSpecification.allowContextTesting() for the Spock test framework. - try { - Class.forName("datadog.context.ContextManager").getMethod("allowTesting").invoke(null); - Class.forName("datadog.context.ContextBinder").getMethod("allowTesting").invoke(null); - } catch (Throwable ignore) { - // don't block testing if context types aren't available - } - } - @BeforeEach public void init() { // If this fails, it's likely the result of another test loading Config before it can be diff --git a/utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTesting.java b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTesting.java new file mode 100644 index 00000000000..ba20511d413 --- /dev/null +++ b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTesting.java @@ -0,0 +1,16 @@ +package datadog.trace.junit.utils.context; + +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; +import org.junit.jupiter.api.extension.ExtendWith; + +/** + * Enables context testing by allowing re-registration of {@code ContextManager} and {@code + * ContextBinder} singletons. + */ +@Retention(RetentionPolicy.RUNTIME) +@Target(ElementType.TYPE) +@ExtendWith(AllowContextTestingExtension.class) +public @interface AllowContextTesting {} diff --git a/utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTestingExtension.java b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTestingExtension.java new file mode 100644 index 00000000000..e07149bdc89 --- /dev/null +++ b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTestingExtension.java @@ -0,0 +1,27 @@ +package datadog.trace.junit.utils.context; + +import de.thetaphi.forbiddenapis.SuppressForbidden; +import org.junit.jupiter.api.extension.BeforeAllCallback; +import org.junit.jupiter.api.extension.ExtensionContext; + +/** + * JUnit 5 extension that allows re-registration of context managers so each test can use a fresh + * tracer. This is needed because {@code ContextManager} and {@code ContextBinder} are singletons + * that normally reject re-registration. + * + *

Auto-registered when using {@link AllowContextTesting}. Can also be used explicitly via + * {@code @ExtendWith(AllowContextTestingExtension.class)}. + */ +@SuppressForbidden +public class AllowContextTestingExtension implements BeforeAllCallback { + + @Override + public void beforeAll(ExtensionContext context) { + try { + Class.forName("datadog.context.ContextManager").getMethod("allowTesting").invoke(null); + Class.forName("datadog.context.ContextBinder").getMethod("allowTesting").invoke(null); + } catch (Throwable ignore) { + // don't block testing if context types aren't available + } + } +} diff --git a/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java b/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java index 028ec3a7bce..44fd75c7327 100644 --- a/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java +++ b/utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java @@ -4,8 +4,8 @@ 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.lang.reflect.Method; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -16,32 +16,23 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.extension.ExtendWith; -@ExtendWith(WithConfigExtension.class) +@ExtendWith({WithConfigExtension.class, AllowContextTestingExtension.class}) @SuppressForbidden public class DDJavaSpecification { private static final long CHECK_TIMEOUT_MS = 3000; - static final String CONTEXT_BINDER = "datadog.context.ContextBinder"; - static final String CONTEXT_MANAGER = "datadog.context.ContextManager"; - - private static Boolean contextTestingAllowed; - protected boolean assertThreadsEachCleanup = true; private static volatile boolean ignoreThreadCleanup; @BeforeAll static void beforeAll() { - allowContextTesting(); assertTrue( EnvironmentVariables.getAll().entrySet().stream() .noneMatch(e -> e.getKey().startsWith("DD_"))); assertTrue( systemPropertiesExceptAllowed().entrySet().stream() .noneMatch(e -> e.getKey().toString().startsWith("dd."))); - assertTrue( - contextTestingAllowed, - "Context not ready for testing. Ensure all test classes extend DDJavaSpecification"); if (getDDThreads().isEmpty()) { ignoreThreadCleanup = false; @@ -52,27 +43,6 @@ static void beforeAll() { } } - static void allowContextTesting() { - if (contextTestingAllowed == null) { - try { - Class binderClass = Class.forName(CONTEXT_BINDER); - Method binderAllowTesting = binderClass.getDeclaredMethod("allowTesting"); - binderAllowTesting.setAccessible(true); - Class managerClass = Class.forName(CONTEXT_MANAGER); - Method managerAllowTesting = managerClass.getDeclaredMethod("allowTesting"); - managerAllowTesting.setAccessible(true); - contextTestingAllowed = - (Boolean) binderAllowTesting.invoke(null) && (Boolean) managerAllowTesting.invoke(null); - } catch (ClassNotFoundException e) { - // don't block testing if these types aren't found (project doesn't use context API) - contextTestingAllowed = - CONTEXT_BINDER.equals(e.getMessage()) || CONTEXT_MANAGER.equals(e.getMessage()); - } catch (Throwable ignore) { - contextTestingAllowed = false; - } - } - } - @AfterAll static void afterAll() { assertTrue( @@ -135,5 +105,4 @@ static void checkThreads() { System.out.println(names); } } - } From 8291efc5675db0b4b53e0e5d8ad4683990d672f4 Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Fri, 10 Apr 2026 10:09:47 +0200 Subject: [PATCH 3/4] feat(junit): Move table tests into junit-utils module --- .../java/datadog/trace/api/DDTraceApiTableTestConverters.java | 2 +- utils/junit-utils/build.gradle.kts | 1 + .../trace/junit/utils/tabletest}/TableTestTypeConverters.java | 2 +- utils/test-utils/build.gradle.kts | 1 - 4 files changed, 3 insertions(+), 3 deletions(-) rename utils/{test-utils/src/main/java/datadog/trace/test/util => junit-utils/src/main/java/datadog/trace/junit/utils/tabletest}/TableTestTypeConverters.java (93%) diff --git a/dd-trace-api/src/test/java/datadog/trace/api/DDTraceApiTableTestConverters.java b/dd-trace-api/src/test/java/datadog/trace/api/DDTraceApiTableTestConverters.java index e277fe8e2e7..bb21c3aac04 100644 --- a/dd-trace-api/src/test/java/datadog/trace/api/DDTraceApiTableTestConverters.java +++ b/dd-trace-api/src/test/java/datadog/trace/api/DDTraceApiTableTestConverters.java @@ -1,6 +1,6 @@ package datadog.trace.api; -import datadog.trace.test.util.TableTestTypeConverters; +import datadog.trace.junit.utils.tabletest.TableTestTypeConverters; import org.tabletest.junit.TypeConverter; /** TableTest converters shared by dd-trace-api test classes for unparsable constants. */ diff --git a/utils/junit-utils/build.gradle.kts b/utils/junit-utils/build.gradle.kts index e26ddcfa1ea..f30b2ebca87 100644 --- a/utils/junit-utils/build.gradle.kts +++ b/utils/junit-utils/build.gradle.kts @@ -11,4 +11,5 @@ dependencies { api(project(":components:environment")) compileOnly(libs.junit.jupiter) + compileOnly(libs.tabletest) } diff --git a/utils/test-utils/src/main/java/datadog/trace/test/util/TableTestTypeConverters.java b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/tabletest/TableTestTypeConverters.java similarity index 93% rename from utils/test-utils/src/main/java/datadog/trace/test/util/TableTestTypeConverters.java rename to utils/junit-utils/src/main/java/datadog/trace/junit/utils/tabletest/TableTestTypeConverters.java index 973408cd3fa..1341f109cbb 100644 --- a/utils/test-utils/src/main/java/datadog/trace/test/util/TableTestTypeConverters.java +++ b/utils/junit-utils/src/main/java/datadog/trace/junit/utils/tabletest/TableTestTypeConverters.java @@ -1,4 +1,4 @@ -package datadog.trace.test.util; +package datadog.trace.junit.utils.tabletest; import org.tabletest.junit.TypeConverter; diff --git a/utils/test-utils/build.gradle.kts b/utils/test-utils/build.gradle.kts index 07b6ccf1e3f..4481b0cbbba 100644 --- a/utils/test-utils/build.gradle.kts +++ b/utils/test-utils/build.gradle.kts @@ -15,7 +15,6 @@ dependencies { api(group = "commons-fileupload", name = "commons-fileupload", version = "1.5") compileOnly(libs.junit.jupiter) - compileOnly(libs.tabletest) compileOnly(libs.logback.core) compileOnly(libs.logback.classic) From 8d909c3d37ac062968e98bac6b060ff0fec6d3c0 Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Fri, 10 Apr 2026 09:59:53 +0200 Subject: [PATCH 4/4] feat(core): Improve span links tests --- .../trace/core/propagation/W3CHttpCodec.java | 11 +-- .../datadog/trace/core/DDSpanLinkTest.java | 74 ++++++++++--------- .../core/propagation/HttpCodecTestHelper.java | 2 +- 3 files changed, 45 insertions(+), 42 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java index cd296c27b8a..52037859cb2 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java @@ -32,11 +32,6 @@ class W3CHttpCodec { private static final Logger log = LoggerFactory.getLogger(W3CHttpCodec.class); - static final String TRACE_PARENT_KEY = "traceparent"; - static final String TRACE_STATE_KEY = "tracestate"; - static final String OT_BAGGAGE_PREFIX = "ot-baggage-"; - private static final String E2E_START_KEY = OT_BAGGAGE_PREFIX + DDTags.TRACE_START_TIME; - private static final int TRACE_PARENT_TID_START = 2 + 1; private static final int TRACE_PARENT_TID_END = TRACE_PARENT_TID_START + 32; private static final int TRACE_PARENT_SID_START = TRACE_PARENT_TID_END + 1; @@ -45,6 +40,12 @@ class W3CHttpCodec { private static final int TRACE_PARENT_FLAGS_SAMPLED = 1; private static final int TRACE_PARENT_LENGTH = TRACE_PARENT_FLAGS_START + 2; + // Package-protected for testing + static final String TRACE_PARENT_KEY = "traceparent"; + static final String TRACE_STATE_KEY = "tracestate"; + static final String OT_BAGGAGE_PREFIX = "ot-baggage-"; + static final String E2E_START_KEY = OT_BAGGAGE_PREFIX + DDTags.TRACE_START_TIME; + private W3CHttpCodec() { // This class should not be created. This also makes code coverage checks happy. } diff --git a/dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java index fd1667da21e..f6000cef186 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java @@ -3,33 +3,35 @@ import static datadog.trace.api.DDTags.SPAN_LINKS; import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.DEFAULT_FLAGS; import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.SAMPLED_FLAG; +import static datadog.trace.bootstrap.instrumentation.api.ContextVisitors.stringValuesMap; import static datadog.trace.bootstrap.instrumentation.api.SpanAttributes.EMPTY; import static datadog.trace.core.propagation.HttpCodecTestHelper.TRACE_PARENT_KEY; import static datadog.trace.core.propagation.HttpCodecTestHelper.TRACE_STATE_KEY; +import static datadog.trace.core.propagation.HttpCodecTestHelper.newW3cHttpCodecExtractor; +import static java.util.Collections.emptyList; +import static java.util.stream.Collectors.toList; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.type.CollectionType; import datadog.trace.api.Config; import datadog.trace.api.DDSpanId; import datadog.trace.api.DDTraceId; import datadog.trace.api.DynamicConfig; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.bootstrap.instrumentation.api.AgentTracer; -import datadog.trace.bootstrap.instrumentation.api.ContextVisitors; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.SpanBuilder; import datadog.trace.bootstrap.instrumentation.api.SpanAttributes; import datadog.trace.bootstrap.instrumentation.api.SpanLink; import datadog.trace.common.writer.ListWriter; import datadog.trace.core.propagation.ExtractedContext; import datadog.trace.core.propagation.HttpCodec; -import datadog.trace.core.propagation.HttpCodecTestHelper; import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.stream.Collectors; import java.util.stream.IntStream; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -41,19 +43,21 @@ class DDSpanLinkTest extends DDCoreJavaSpecification { private static final int SPAN_LINK_TAG_MAX_LENGTH = 25_000; private static final ObjectMapper JSON_MAPPER = new ObjectMapper(); + private static final CollectionType SPAN_LINK_LIST_TYPE = + JSON_MAPPER.getTypeFactory().constructCollectionType(List.class, SpanLinkAsTag.class); private ListWriter writer; private CoreTracer tracer; @BeforeEach void setup() { - writer = new ListWriter(); - tracer = tracerBuilder().writer(writer).build(); + this.writer = new ListWriter(); + this.tracer = tracerBuilder().writer(this.writer).build(); } @AfterEach void cleanupTest() { - writer.clear(); + this.writer.clear(); } @TableTest({ @@ -65,15 +69,15 @@ void createSpanLinkFromExtractedContext(boolean sampled, String traceFlags, Stri String traceId = "11223344556677889900aabbccddeeff"; String spanId = "123456789abcdef0"; String traceState = "dd=s:" + sample + ";o:some;t.dm:-4"; + Map headers = new HashMap<>(); headers.put(TRACE_PARENT_KEY.toUpperCase(), "00-" + traceId + "-" + spanId + "-" + traceFlags); headers.put(TRACE_STATE_KEY.toUpperCase(), traceState); HttpCodec.Extractor extractor = - HttpCodecTestHelper.W3CHttpCodecNewExtractor( + newW3cHttpCodecExtractor( Config.get(), () -> DynamicConfig.create().apply().captureTraceConfig()); - ExtractedContext context = - (ExtractedContext) extractor.extract(headers, ContextVisitors.stringValuesMap()); + ExtractedContext context = (ExtractedContext) extractor.extract(headers, stringValuesMap()); SpanLink link = DDSpanLink.from(context); assertEquals(DDTraceId.fromHex(traceId), link.traceId()); @@ -85,19 +89,18 @@ void createSpanLinkFromExtractedContext(boolean sampled, String traceFlags, Stri @Test void testSpanLinkEncodingTagMaxSize() throws Exception { int tooManyLinkCount = 300; - AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); + SpanBuilder builder = tracer.buildSpan("test", "operation"); List links = IntStream.range(0, tooManyLinkCount) .mapToObj(this::createLink) - .collect(Collectors.toList()); - - for (SpanLink link : links) { - builder.withLink(link); - } + .peek(builder::withLink) + .collect(toList()); AgentSpan span = builder.start(); span.finish(); - writer.waitForTraces(1); - String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); + this.writer.waitForTraces(1); + + assertEquals(1, this.writer.get(0).size()); + String spanLinksTag = (String) this.writer.get(0).get(0).getTag(SPAN_LINKS); List decodedSpanLinks = deserializeSpanLinks(spanLinksTag); assertTrue(spanLinksTag.length() < SPAN_LINK_TAG_MAX_LENGTH); @@ -112,7 +115,6 @@ void testSpanLinkEncodingTagMaxSize() throws Exception { @Test void testSpanLinksEncodingOmittedEmptyKeys() throws Exception { - AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); SpanLink link = new DDSpanLink( DDTraceId.fromHex("11223344556677889900aabbccddeeff"), @@ -120,12 +122,11 @@ void testSpanLinksEncodingOmittedEmptyKeys() throws Exception { DEFAULT_FLAGS, "", EMPTY); + this.tracer.buildSpan("test", "operation").withLink(link).start().finish(); + this.writer.waitForTraces(1); - AgentSpan span = builder.withLink(link).start(); - span.finish(); - writer.waitForTraces(1); + assertEquals(1, this.writer.get(0).size()); String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); - assertEquals( "[{\"span_id\":\"123456789abcdef0\",\"trace_id\":\"11223344556677889900aabbccddeeff\"}]", spanLinksTag); @@ -140,7 +141,7 @@ void testSpanLinksEncodingOmittedEmptyKeys() throws Exception { }) @ParameterizedTest(name = "add span link at any time [{index}]") void addSpanLinkAtAnyTime(boolean beforeStart, boolean afterStart) throws Exception { - AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); + SpanBuilder builder = this.tracer.buildSpan("test", "operation"); List links = new ArrayList<>(); if (beforeStart) { @@ -155,12 +156,11 @@ void addSpanLinkAtAnyTime(boolean beforeStart, boolean afterStart) throws Except links.add(link); } span.finish(); - writer.waitForTraces(1); - String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); - List decodedSpanLinks = - spanLinksTag == null - ? java.util.Collections.emptyList() - : deserializeSpanLinks(spanLinksTag); + this.writer.waitForTraces(1); + + assertEquals(1, this.writer.get(0).size()); + String spanLinksTag = (String) this.writer.get(0).get(0).getTag(SPAN_LINKS); + List decodedSpanLinks = deserializeSpanLinks(spanLinksTag); int expectedLinkCount = (beforeStart ? 1 : 0) + (afterStart ? 1 : 0); assertEquals(expectedLinkCount, decodedSpanLinks.size()); @@ -171,14 +171,15 @@ void addSpanLinkAtAnyTime(boolean beforeStart, boolean afterStart) throws Except @Test void filterNullLinks() throws Exception { - AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation"); + SpanBuilder builder = this.tracer.buildSpan("test", "operation"); AgentSpan span = builder.withLink(null).start(); span.addLink(null); span.finish(); - writer.waitForTraces(1); - String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS); + this.writer.waitForTraces(1); + assertEquals(1, this.writer.get(0).size()); + String spanLinksTag = (String) this.writer.get(0).get(0).getTag(SPAN_LINKS); assertNull(spanLinksTag); } @@ -215,9 +216,10 @@ private void assertLink(SpanLink expected, SpanLinkAsTag actual) { } static List deserializeSpanLinks(String json) throws IOException { - return JSON_MAPPER.readValue( - json, - JSON_MAPPER.getTypeFactory().constructCollectionType(List.class, SpanLinkAsTag.class)); + if (json == null) { + return emptyList(); + } + return JSON_MAPPER.readValue(json, SPAN_LINK_LIST_TYPE); } static class SpanLinkAsTag { diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java index c856cdaf8dc..00e0d460572 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java @@ -10,7 +10,7 @@ public class HttpCodecTestHelper { public static final String TRACE_PARENT_KEY = W3CHttpCodec.TRACE_PARENT_KEY; public static final String TRACE_STATE_KEY = W3CHttpCodec.TRACE_STATE_KEY; - public static HttpCodec.Extractor W3CHttpCodecNewExtractor( + public static HttpCodec.Extractor newW3cHttpCodecExtractor( Config config, Supplier traceConfigSupplier) { return W3CHttpCodec.newExtractor(config, traceConfigSupplier); }