Skip to content

Commit c13e821

Browse files
PerfectSlayerPerfectSlayer
authored
Improve Config JUnit extension (#11094)
feat(junit): Fix @WithConfig to be applied at any level fix(junit): Only rebuild config once fix(junit): Apply class level configuration at @beforeeach level So child classes can benefits from class level configs and we can configure the tracer for example feat(test): Extract config cleanup assertions into dedicated JUnit extension Move DD_* environment variable and dd.* system property validation logic from DDJavaSpecification into a reusable CleanConfigStateExtension. Improves error reporting with detailed leak information. feat(junit): Improve generated code readability Co-authored-by: bruce.bujon <bruce.bujon@datadoghq.com>
1 parent c99e05c commit c13e821

File tree

5 files changed

+159
-59
lines changed

5 files changed

+159
-59
lines changed

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package datadog.trace.junit.utils.config;
22

3-
import java.lang.annotation.ElementType;
3+
import static java.lang.annotation.ElementType.METHOD;
4+
import static java.lang.annotation.ElementType.TYPE;
5+
import static java.lang.annotation.RetentionPolicy.RUNTIME;
6+
47
import java.lang.annotation.Repeatable;
58
import java.lang.annotation.Retention;
6-
import java.lang.annotation.RetentionPolicy;
79
import java.lang.annotation.Target;
810
import org.junit.jupiter.api.extension.ExtendWith;
911

@@ -27,8 +29,8 @@
2729
* }
2830
* }</pre>
2931
*/
30-
@Retention(RetentionPolicy.RUNTIME)
31-
@Target({ElementType.TYPE, ElementType.METHOD})
32+
@Retention(RUNTIME)
33+
@Target({TYPE, METHOD})
3234
@Repeatable(WithConfigs.class)
3335
@ExtendWith(WithConfigExtension.class)
3436
public @interface WithConfig {

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

Lines changed: 62 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.lang.reflect.Constructor;
2020
import java.lang.reflect.Field;
2121
import java.lang.reflect.Modifier;
22+
import java.util.ArrayList;
2223
import java.util.HashMap;
2324
import java.util.List;
2425
import java.util.Map;
@@ -62,6 +63,7 @@ public class WithConfigExtension
6263
private static Field configInstanceField;
6364
private static Constructor<?> configConstructor;
6465

66+
private static volatile boolean configTransformerInstalled = false;
6567
private static volatile boolean isConfigInstanceModifiable = false;
6668
private static volatile boolean configModificationFailed = false;
6769

@@ -73,22 +75,42 @@ public class WithConfigExtension
7375

7476
@Override
7577
public void beforeAll(ExtensionContext context) {
76-
installConfigTransformer();
78+
/*
79+
* Patch config classes to make them modifiable.
80+
*/
81+
// Install config transformer error listener
82+
if (!configTransformerInstalled) {
83+
installConfigTransformer();
84+
configTransformerInstalled = true;
85+
}
86+
// Make config instance modifiable
7787
makeConfigInstanceModifiable();
88+
// Verify that config class transformation succeeded
7889
assertFalse(configModificationFailed, "Config class modification failed");
90+
if (isConfigInstanceModifiable) {
91+
checkConfigTransformation();
92+
}
93+
/*
94+
* Back up config and apply class-level config values.
95+
*/
7996
if (originalSystemProperties == null) {
8097
saveProperties();
8198
}
99+
// Apply class-level @WithConfig so config is available before @BeforeAll methods
100+
applyClassLevelConfig(context);
101+
if (isConfigInstanceModifiable) {
102+
rebuildConfig();
103+
}
82104
}
83105

84106
@Override
85107
public void beforeEach(ExtensionContext context) {
86108
restoreProperties();
87109
environmentVariables.clear();
110+
applyDeclaredConfig(context);
88111
if (isConfigInstanceModifiable) {
89112
rebuildConfig();
90113
}
91-
applyDeclaredConfig(context);
92114
}
93115

94116
@Override
@@ -108,14 +130,29 @@ public void afterAll(ExtensionContext context) {
108130
}
109131
}
110132

111-
private void applyDeclaredConfig(ExtensionContext context) {
112-
// Class-level @WithConfig annotations (supports composed/meta-annotations)
113-
List<WithConfig> classConfigs =
114-
AnnotationSupport.findRepeatableAnnotations(
115-
context.getRequiredTestClass(), WithConfig.class);
116-
for (WithConfig cfg : classConfigs) {
117-
applyConfig(cfg);
133+
private static void applyDeclaredConfig(ExtensionContext context) {
134+
applyClassLevelConfig(context);
135+
applyMethodLevelConfig(context);
136+
}
137+
138+
private static void applyClassLevelConfig(ExtensionContext context) {
139+
// Walk the entire class hierarchy so annotations on superclasses and apply topmost first, then
140+
// subclass overrides.
141+
Class<?> testClass = context.getRequiredTestClass();
142+
List<Class<?>> hierarchy = new ArrayList<>();
143+
for (Class<?> cls = testClass; cls != null; cls = cls.getSuperclass()) {
144+
hierarchy.add(cls);
145+
}
146+
for (int i = hierarchy.size() - 1; i >= 0; i--) {
147+
List<WithConfig> classConfigs =
148+
AnnotationSupport.findRepeatableAnnotations(hierarchy.get(i), WithConfig.class);
149+
for (WithConfig cfg : classConfigs) {
150+
applyConfig(cfg);
151+
}
118152
}
153+
}
154+
155+
private static void applyMethodLevelConfig(ExtensionContext context) {
119156
// Method-level @WithConfig annotations (supports composed/meta-annotations)
120157
context
121158
.getTestMethod()
@@ -131,12 +168,22 @@ private void applyDeclaredConfig(ExtensionContext context) {
131168

132169
private static void applyConfig(WithConfig cfg) {
133170
if (cfg.env()) {
134-
injectEnvConfig(cfg.key(), cfg.value(), cfg.addPrefix());
171+
setEnvVariable(cfg.key(), cfg.value(), cfg.addPrefix());
135172
} else {
136-
injectSysConfig(cfg.key(), cfg.value(), cfg.addPrefix());
173+
setSysProperty(cfg.key(), cfg.value(), cfg.addPrefix());
137174
}
138175
}
139176

177+
private static void setSysProperty(String name, String value, boolean addPrefix) {
178+
String prefixedName = addPrefix && !name.startsWith("dd.") ? "dd." + name : name;
179+
System.setProperty(prefixedName, value);
180+
}
181+
182+
private static void setEnvVariable(String name, String value, boolean addPrefix) {
183+
String prefixedName = addPrefix && !name.startsWith("DD_") ? "DD_" + name : name;
184+
environmentVariables.set(prefixedName, value);
185+
}
186+
140187
// endregion
141188

142189
// region Public static API for imperative config injection
@@ -146,9 +193,7 @@ public static void injectSysConfig(String name, String value) {
146193
}
147194

148195
public static void injectSysConfig(String name, String value, boolean addPrefix) {
149-
checkConfigTransformation();
150-
String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name;
151-
System.setProperty(prefixedName, value);
196+
setSysProperty(name, value, addPrefix);
152197
rebuildConfig();
153198
}
154199

@@ -157,8 +202,7 @@ public static void removeSysConfig(String name) {
157202
}
158203

159204
public static void removeSysConfig(String name, boolean addPrefix) {
160-
checkConfigTransformation();
161-
String prefixedName = name.startsWith("dd.") || !addPrefix ? name : "dd." + name;
205+
String prefixedName = addPrefix && !name.startsWith("dd.") ? "dd." + name : name;
162206
System.clearProperty(prefixedName);
163207
rebuildConfig();
164208
}
@@ -168,9 +212,7 @@ public static void injectEnvConfig(String name, String value) {
168212
}
169213

170214
public static void injectEnvConfig(String name, String value, boolean addPrefix) {
171-
checkConfigTransformation();
172-
String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name;
173-
environmentVariables.set(prefixedName, value);
215+
setEnvVariable(name, value, addPrefix);
174216
rebuildConfig();
175217
}
176218

@@ -179,8 +221,7 @@ public static void removeEnvConfig(String name) {
179221
}
180222

181223
public static void removeEnvConfig(String name, boolean addPrefix) {
182-
checkConfigTransformation();
183-
String prefixedName = name.startsWith("DD_") || !addPrefix ? name : "DD_" + name;
224+
String prefixedName = addPrefix && !name.startsWith("DD_") ? "DD_" + name : name;
184225
environmentVariables.removePrefixed(prefixedName);
185226
rebuildConfig();
186227
}
@@ -245,7 +286,6 @@ static void makeConfigInstanceModifiable() {
245286

246287
private static void rebuildConfig() {
247288
synchronized (WithConfigExtension.class) {
248-
checkConfigTransformation();
249289
try {
250290
Object newInstConfig = instConfigConstructor.newInstance();
251291
instConfigInstanceField.set(null, newInstConfig);

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
11
package datadog.trace.junit.utils.config;
22

3-
import java.lang.annotation.ElementType;
3+
import static java.lang.annotation.ElementType.METHOD;
4+
import static java.lang.annotation.ElementType.TYPE;
5+
import static java.lang.annotation.RetentionPolicy.RUNTIME;
6+
47
import java.lang.annotation.Retention;
5-
import java.lang.annotation.RetentionPolicy;
68
import java.lang.annotation.Target;
79
import org.junit.jupiter.api.extension.ExtendWith;
810

911
/** Container annotation for repeatable {@link WithConfig}. */
10-
@Retention(RetentionPolicy.RUNTIME)
11-
@Target({ElementType.TYPE, ElementType.METHOD})
12+
@Retention(RUNTIME)
13+
@Target({TYPE, METHOD})
1214
@ExtendWith(WithConfigExtension.class)
1315
public @interface WithConfigs {
1416
WithConfig[] value();

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package datadog.trace.test.util;
2+
3+
import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
4+
5+
import datadog.environment.EnvironmentVariables;
6+
import de.thetaphi.forbiddenapis.SuppressForbidden;
7+
import java.util.Arrays;
8+
import java.util.List;
9+
import java.util.Map;
10+
import java.util.Map.Entry;
11+
import java.util.TreeMap;
12+
import java.util.function.Predicate;
13+
import java.util.stream.Collectors;
14+
import org.junit.jupiter.api.extension.AfterAllCallback;
15+
import org.junit.jupiter.api.extension.BeforeAllCallback;
16+
import org.junit.jupiter.api.extension.ExtensionContext;
17+
18+
/**
19+
* Asserts that no {@code DD_*} environment variable and no {@code dd.*} system property (minus a
20+
* small allowlist) is set around a test class.
21+
*/
22+
@SuppressForbidden
23+
public class CleanConfigStateExtension implements BeforeAllCallback, AfterAllCallback {
24+
25+
private static final List<String> ALLOWED_SYS_PROPS =
26+
Arrays.asList(
27+
"dd.appsec.enabled", "dd.iast.enabled", "dd.integration.grizzly-filterchain.enabled");
28+
29+
private static final Predicate<String> DATADOG_ENV_VAR_FILTER = k -> k.startsWith("DD_");
30+
private static final Predicate<Object> DATADOG_SYS_PROPERTIES_FILTER =
31+
o -> {
32+
String key = (String) o;
33+
return key.startsWith("DD_") && !ALLOWED_SYS_PROPS.contains(key);
34+
};
35+
36+
@Override
37+
public void beforeAll(ExtensionContext context) {
38+
assertClean("before");
39+
}
40+
41+
@Override
42+
public void afterAll(ExtensionContext context) {
43+
assertClean("after");
44+
}
45+
46+
private static void assertClean(String phase) {
47+
Map<String, String> leakedEnv =
48+
filterMap(EnvironmentVariables.getAll(), DATADOG_ENV_VAR_FILTER);
49+
Map<Object, Object> leakedSys =
50+
filterMap(System.getProperties(), DATADOG_SYS_PROPERTIES_FILTER);
51+
if (!leakedEnv.isEmpty() || !leakedSys.isEmpty()) {
52+
assertionFailure()
53+
.message("Leaked Datadog configuration detected " + phase + " test class")
54+
.reason(formatLeaks(leakedEnv, leakedSys))
55+
.buildAndThrow();
56+
}
57+
}
58+
59+
private static <T> Map<T, T> filterMap(Map<T, T> map, Predicate<T> keyFilter) {
60+
return map.entrySet().stream()
61+
.filter(e -> keyFilter.test(e.getKey()))
62+
.collect(Collectors.toMap(Entry::getKey, Entry::getValue, (a, b) -> a, TreeMap::new));
63+
}
64+
65+
private static String formatLeaks(Map<String, String> env, Map<Object, Object> sys) {
66+
StringBuilder sb = new StringBuilder();
67+
if (!env.isEmpty()) {
68+
sb.append("environment variables:");
69+
env.forEach((k, v) -> sb.append("\n ").append(k).append('=').append(v));
70+
}
71+
if (!sys.isEmpty()) {
72+
if (sb.length() > 0) {
73+
sb.append('\n');
74+
}
75+
sb.append("system properties:");
76+
sys.forEach((k, v) -> sb.append("\n ").append(k).append('=').append(v));
77+
}
78+
return sb.toString();
79+
}
80+
}

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

Lines changed: 5 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,21 @@
11
package datadog.trace.test.util;
22

3-
import static org.junit.jupiter.api.Assertions.assertTrue;
4-
5-
import datadog.environment.EnvironmentVariables;
63
import datadog.trace.junit.utils.config.WithConfigExtension;
74
import datadog.trace.junit.utils.context.AllowContextTestingExtension;
85
import de.thetaphi.forbiddenapis.SuppressForbidden;
9-
import java.util.Arrays;
106
import java.util.List;
11-
import java.util.Map;
127
import java.util.Set;
138
import java.util.stream.Collectors;
149
import org.junit.jupiter.api.AfterAll;
1510
import org.junit.jupiter.api.AfterEach;
1611
import org.junit.jupiter.api.BeforeAll;
1712
import org.junit.jupiter.api.extension.ExtendWith;
1813

19-
@ExtendWith({WithConfigExtension.class, AllowContextTestingExtension.class})
14+
@ExtendWith({
15+
CleanConfigStateExtension.class,
16+
WithConfigExtension.class,
17+
AllowContextTestingExtension.class
18+
})
2019
@SuppressForbidden
2120
public class DDJavaSpecification {
2221

@@ -27,13 +26,6 @@ public class DDJavaSpecification {
2726

2827
@BeforeAll
2928
static void beforeAll() {
30-
assertTrue(
31-
EnvironmentVariables.getAll().entrySet().stream()
32-
.noneMatch(e -> e.getKey().startsWith("DD_")));
33-
assertTrue(
34-
systemPropertiesExceptAllowed().entrySet().stream()
35-
.noneMatch(e -> e.getKey().toString().startsWith("dd.")));
36-
3729
if (getDDThreads().isEmpty()) {
3830
ignoreThreadCleanup = false;
3931
} else {
@@ -45,25 +37,9 @@ static void beforeAll() {
4537

4638
@AfterAll
4739
static void afterAll() {
48-
assertTrue(
49-
EnvironmentVariables.getAll().entrySet().stream()
50-
.noneMatch(e -> e.getKey().startsWith("DD_")));
51-
assertTrue(
52-
systemPropertiesExceptAllowed().entrySet().stream()
53-
.noneMatch(e -> e.getKey().toString().startsWith("dd.")));
54-
5540
checkThreads();
5641
}
5742

58-
private static Map<Object, Object> systemPropertiesExceptAllowed() {
59-
List<String> allowlist =
60-
Arrays.asList(
61-
"dd.appsec.enabled", "dd.iast.enabled", "dd.integration.grizzly-filterchain.enabled");
62-
return System.getProperties().entrySet().stream()
63-
.filter(e -> !allowlist.contains(String.valueOf(e.getKey())))
64-
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
65-
}
66-
6743
@AfterEach
6844
void cleanup() {
6945
if (assertThreadsEachCleanup) {

0 commit comments

Comments
 (0)