1414import datadog .trace .api .Config ;
1515import datadog .trace .api .IdGenerationStrategy ;
1616import datadog .trace .bootstrap .instrumentation .api .AgentSpan ;
17- import datadog .trace .bootstrap .instrumentation .api .AgentTracer ;
17+ import datadog .trace .bootstrap .instrumentation .api .AgentTracer . TracerAPI ;
1818import datadog .trace .common .writer .ListWriter ;
1919import datadog .trace .core .CoreTracer ;
2020import datadog .trace .core .DDSpan ;
3030import java .util .function .Function ;
3131import java .util .function .Predicate ;
3232import net .bytebuddy .agent .ByteBuddyAgent ;
33+ import org .junit .jupiter .api .AfterAll ;
3334import org .junit .jupiter .api .AfterEach ;
35+ import org .junit .jupiter .api .BeforeAll ;
3436import org .junit .jupiter .api .BeforeEach ;
3537import org .junit .jupiter .api .extension .ExtendWith ;
3638import org .opentest4j .AssertionFailedError ;
3739
3840/**
39- * This class is an experimental base to run instrumentation tests using JUnit Jupiter. It is still
40- * early development, and the overall API is expected to change to leverage its extension model. The
41- * current implementation is inspired and kept close to it Groovy / Spock counterpart, the {@code
42- * InstrumentationSpecification}.
41+ * Base class for instrumentation tests using JUnit Jupiter.
42+ *
43+ * <p>It is still early development, and the overall API might change to leverage its extension
44+ * model. The current implementation is inspired and kept close to its Groovy / Spock counterpart,
45+ * the {@code InstrumentationSpecification}.
46+ *
47+ * <ul>
48+ * <li>{@code @BeforeAll}: Installs the agent and creates a shared tracer
49+ * <li>{@code @BeforeEach}: Flushes and resets the writer
50+ * <li>{@code @AfterEach}: Flushes the tracer
51+ * <li>{@code @AfterAll}: Closes the tracer and removes the agent transformer
52+ * </ul>
4353 */
4454@ ExtendWith ({TestClassShadowingExtension .class , AllowContextTestingExtension .class })
4555public abstract class AbstractInstrumentationTest {
4656 static final Instrumentation INSTRUMENTATION = ByteBuddyAgent .getInstrumentation ();
4757
4858 static final long TIMEOUT_MILLIS = TimeUnit .SECONDS .toMillis (20 );
4959
50- protected AgentTracer . TracerAPI tracer ;
60+ protected static final InstrumentationTestConfig testConfig = new InstrumentationTestConfig () ;
5161
52- protected ListWriter writer ;
62+ protected static TracerAPI tracer ;
63+ protected static ListWriter writer ;
64+ private static ClassFileTransformer activeTransformer ;
65+ private static ClassFileTransformerListener transformerListener ;
5366
54- protected ClassFileTransformer activeTransformer ;
55- protected ClassFileTransformerListener transformerLister ;
56-
57- @ BeforeEach
58- public void init () {
67+ @ BeforeAll
68+ static void initAll () {
5969 // If this fails, it's likely the result of another test loading Config before it can be
6070 // injected into the bootstrap classpath.
61- // If one test extends AgentTestRunner in a module, all tests must extend
6271 assertNull (Config .class .getClassLoader (), "Config must load on the bootstrap classpath." );
6372
64- // Initialize test tracer
65- this .writer = new ListWriter ();
66- // Initialize test tracer
67- CoreTracer tracer =
73+ // Create shared test writer and tracer
74+ writer = new ListWriter ();
75+ CoreTracer coreTracer =
6876 CoreTracer .builder ()
69- .writer (this . writer )
70- .idGenerationStrategy (IdGenerationStrategy .fromName (idGenerationStrategyName () ))
71- .strictTraceWrites (useStrictTraceWrites () )
77+ .writer (writer )
78+ .idGenerationStrategy (IdGenerationStrategy .fromName (testConfig . idGenerationStrategy ))
79+ .strictTraceWrites (testConfig . strictTraceWrites )
7280 .build ();
73- TracerInstaller .forceInstallGlobalTracer (tracer );
74- this . tracer = tracer ;
81+ TracerInstaller .forceInstallGlobalTracer (coreTracer );
82+ tracer = coreTracer ;
7583
7684 ClassInjector .enableClassInjection (INSTRUMENTATION );
7785
@@ -85,33 +93,43 @@ public void init() {
8593 .iterator ()
8694 .hasNext (),
8795 "No instrumentation found" );
88- this . transformerLister = new ClassFileTransformerListener ();
89- this . activeTransformer =
96+ transformerListener = new ClassFileTransformerListener ();
97+ activeTransformer =
9098 AgentInstaller .installBytebuddyAgent (
91- INSTRUMENTATION , true , AgentInstaller .getEnabledSystems (), this .transformerLister );
92- }
93-
94- protected String idGenerationStrategyName () {
95- return "SEQUENTIAL" ;
99+ INSTRUMENTATION , true , AgentInstaller .getEnabledSystems (), transformerListener );
96100 }
97101
98- private boolean useStrictTraceWrites () {
99- return true ;
102+ @ BeforeEach
103+ public void init () {
104+ tracer .flush ();
105+ writer .start ();
100106 }
101107
102108 @ AfterEach
103109 public void tearDown () {
104- this .tracer .close ();
105- this .writer .close ();
106- if (this .activeTransformer != null ) {
107- INSTRUMENTATION .removeTransformer (this .activeTransformer );
108- this .activeTransformer = null ;
109- }
110+ tracer .flush ();
111+ }
110112
111- // All cleanups should happen before these assertions.
113+ @ AfterAll
114+ static void tearDownAll () {
115+ if (tracer != null ) {
116+ tracer .close ();
117+ tracer = null ;
118+ }
119+ if (writer != null ) {
120+ writer .close ();
121+ writer = null ;
122+ }
123+ if (activeTransformer != null ) {
124+ INSTRUMENTATION .removeTransformer (activeTransformer );
125+ activeTransformer = null ;
126+ }
127+ // All cleanups should happen before this verify call.
112128 // If not, a failing assertion may prevent cleanup
113- this .transformerLister .verify ();
114- this .transformerLister = null ;
129+ if (transformerListener != null ) {
130+ transformerListener .verify ();
131+ transformerListener = null ;
132+ }
115133 }
116134
117135 /**
@@ -134,11 +152,11 @@ protected void assertTraces(
134152 TraceMatcher ... matchers ) {
135153 int expectedTraceCount = matchers .length ;
136154 try {
137- this . writer .waitForTraces (expectedTraceCount );
155+ writer .waitForTraces (expectedTraceCount );
138156 } catch (InterruptedException | TimeoutException e ) {
139157 throw new AssertionFailedError ("Timeout while waiting for traces" , e );
140158 }
141- TraceAssertions .assertTraces (this . writer , options , matchers );
159+ TraceAssertions .assertTraces (writer , options , matchers );
142160 }
143161
144162 /**
@@ -149,7 +167,7 @@ protected void assertTraces(
149167 */
150168 protected void blockUntilTracesMatch (Predicate <List <List <DDSpan >>> predicate ) {
151169 long deadline = System .currentTimeMillis () + TIMEOUT_MILLIS ;
152- while (!predicate .test (this . writer )) {
170+ while (!predicate .test (writer )) {
153171 if (System .currentTimeMillis () > deadline ) {
154172 throw new RuntimeException (new TimeoutException ("Timed out waiting for traces/spans." ));
155173 }
@@ -161,8 +179,8 @@ protected void blockUntilTracesMatch(Predicate<List<List<DDSpan>>> predicate) {
161179 }
162180 }
163181
164- protected void blockUntilChildSpansFinished (final int numberOfSpans ) {
165- blockUntilChildSpansFinished (this . tracer .activeSpan (), numberOfSpans );
182+ protected void blockUntilChildSpansFinished (int numberOfSpans ) {
183+ blockUntilChildSpansFinished (tracer .activeSpan (), numberOfSpans );
166184 }
167185
168186 static void blockUntilChildSpansFinished (AgentSpan span , int numberOfSpans ) {
@@ -190,4 +208,20 @@ static void blockUntilChildSpansFinished(AgentSpan span, int numberOfSpans) {
190208 }
191209 }
192210 }
211+
212+ /** Configuration for {@link AbstractInstrumentationTest}. */
213+ protected static class InstrumentationTestConfig {
214+ private String idGenerationStrategy = "SEQUENTIAL" ;
215+ private boolean strictTraceWrites = true ;
216+
217+ public InstrumentationTestConfig idGenerationStrategy (String strategy ) {
218+ this .idGenerationStrategy = strategy ;
219+ return this ;
220+ }
221+
222+ public InstrumentationTestConfig strictTraceWrites (boolean strict ) {
223+ this .strictTraceWrites = strict ;
224+ return this ;
225+ }
226+ }
193227}
0 commit comments