[test optimization] Activate test suite span in beforeAll/afterAll hooks in jest (#7770)

juan-fernandez · web-flow · commit 6c8e945ab222 · 2026-03-16T10:55:58.000+01:00
diff --git a/eslint.config.mjs b/eslint.config.mjs
@@ -673,6 +673,7 @@ export default [
     languageOptions: {
       globals: {
         afterAll: 'readonly',
+        beforeAll: 'readonly',
         expect: 'readonly',
         jest: 'readonly',
       },
diff --git a/integration-tests/ci-visibility/test-suite-custom-tags/suite-custom-tags.js b/integration-tests/ci-visibility/test-suite-custom-tags/suite-custom-tags.js
@@ -0,0 +1,26 @@
+'use strict'
+
+const tracer = require('dd-trace')
+const assert = require('assert')
+
+const sum = require('../test/sum')
+
+describe('test optimization', () => {
+  beforeAll(() => {
+    const suiteSpan = tracer.scope().active()
+    suiteSpan.setTag('suite.beforeAll', 'true')
+
+    tracer.trace('beforeAll.setup', () => {})
+  })
+
+  afterAll(() => {
+    const suiteSpan = tracer.scope().active()
+    suiteSpan.setTag('suite.afterAll', 'true')
+
+    tracer.trace('afterAll.teardown', () => {})
+  })
+
+  it('can report tests', () => {
+    assert.strictEqual(sum(1, 2), 3)
+  })
+})
diff --git a/integration-tests/jest/jest.spec.js b/integration-tests/jest/jest.spec.js
@@ -6386,6 +6386,54 @@ describe(`jest@${JEST_VERSION} commonJS`, () => {
         }).catch(done)
       })
     })
+
+    it('does detect custom tags on test suites from beforeAll and afterAll hooks', (done) => {
+      const eventsPromise = receiver
+        .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/citestcycle'), payloads => {
+          const events = payloads.flatMap(({ payload }) => payload.events)
+          const testSuite = events.find(event => event.type === 'test_suite_end').content
+
+          assertObjectContains(testSuite, {
+            meta: {
+              'suite.beforeAll': 'true',
+              'suite.afterAll': 'true',
+            },
+          })
+
+          const suiteSpanId = testSuite.test_suite_id.toString()
+          const sessionTraceId = testSuite.test_session_id.toString()
+
+          // Spans created in beforeAll/afterAll appear as 'span' events and are children of the test suite span
+          const spans = events.filter(event => event.type === 'span').map(event => event.content)
+          const beforeAllSpan = spans.find(span => span.resource === 'beforeAll.setup')
+          const afterAllSpan = spans.find(span => span.resource === 'afterAll.teardown')
+
+          assert.ok(beforeAllSpan)
+          assert.strictEqual(beforeAllSpan.parent_id.toString(), suiteSpanId)
+          assert.strictEqual(beforeAllSpan.trace_id.toString(), sessionTraceId)
+
+          assert.ok(afterAllSpan)
+          assert.strictEqual(afterAllSpan.parent_id.toString(), suiteSpanId)
+          assert.strictEqual(afterAllSpan.trace_id.toString(), sessionTraceId)
+        })
+
+      childProcess = exec(
+        runTestsCommand,
+        {
+          cwd,
+          env: {
+            ...getCiVisAgentlessConfig(receiver.port),
+            TESTS_TO_RUN: 'ci-visibility/test-suite-custom-tags',
+          },
+        }
+      )
+
+      childProcess.on('exit', () => {
+        eventsPromise.then(() => {
+          done()
+        }).catch(done)
+      })
+    })
   })
 
   context('impacted tests', () => {
diff --git a/packages/datadog-instrumentations/src/jest.js b/packages/datadog-instrumentations/src/jest.js
@@ -47,6 +47,7 @@ const testSkippedCh = channel('ci:jest:test:skip')
 const testFinishCh = channel('ci:jest:test:finish')
 const testErrCh = channel('ci:jest:test:err')
 const testFnCh = channel('ci:jest:test:fn')
+const testSuiteHookFnCh = channel('ci:jest:test-suite:hook:fn')
 
 const skippableSuitesCh = channel('ci:jest:test-suite:skippable')
 const libraryConfigurationCh = channel('ci:jest:library-configuration')
@@ -505,6 +506,19 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
         })
       }
 
+      if (event.name === 'hook_start' && (event.hook.type === 'beforeAll' || event.hook.type === 'afterAll')) {
+        const ctx = { testSuiteAbsolutePath: this.testSuiteAbsolutePath }
+        let hookFn = event.hook.fn
+        if (originalHookFns.has(event.hook)) {
+          hookFn = originalHookFns.get(event.hook)
+        } else {
+          originalHookFns.set(event.hook, hookFn)
+        }
+        event.hook.fn = shimmer.wrapFunction(hookFn, hookFn => function () {
+          return testSuiteHookFnCh.runStores(ctx, () => hookFn.apply(this, arguments))
+        })
+      }
+
       if (event.name === 'add_test') {
         if (event.failing) {
           return
diff --git a/packages/datadog-plugin-jest/src/index.js b/packages/datadog-plugin-jest/src/index.js
@@ -386,6 +386,12 @@ class JestPlugin extends CiPlugin {
       return ctx.currentStore
     })
 
+    this.addBind('ci:jest:test-suite:hook:fn', (ctx) => {
+      const testSuiteSpan = this.testSuiteSpanPerTestSuiteAbsolutePath.get(ctx.testSuiteAbsolutePath)
+      const store = storage('legacy').getStore()
+      return { ...store, span: testSuiteSpan }
+    })
+
     this.addSub('ci:jest:test:finish', ({
       span,
       status,
diff --git a/packages/dd-trace/src/encode/agentless-ci-visibility.js b/packages/dd-trace/src/encode/agentless-ci-visibility.js
@@ -265,14 +265,7 @@ class AgentlessCiVisibilityEncoder extends AgentEncoder {
     }
     const startTime = Date.now()
 
-    const rawEvents = trace.map(formatSpan)
-
-    const testSessionEvents = rawEvents.filter(
-      event => event.type === 'test_session_end' || event.type === 'test_suite_end' || event.type === 'test_module_end'
-    )
-
-    const isTestSessionTrace = !!testSessionEvents.length
-    const events = isTestSessionTrace ? testSessionEvents : rawEvents
+    const events = trace.map(formatSpan)
 
     this._eventCount += events.length
 
diff --git a/packages/dd-trace/test/encode/agentless-ci-visibility.spec.js b/packages/dd-trace/test/encode/agentless-ci-visibility.spec.js
@@ -224,8 +224,8 @@ describe('agentless-ci-visibility-encode', () => {
     })
   })
 
-  it('should not encode events other than sessions and suites if the trace is a test session', () => {
-    const traceToFilter = [
+  it('should encode all events including non-test spans alongside test sessions', () => {
+    const traceWithMixedSpans = [
       {
         trace_id: id('1234abcd1234abcd'),
         span_id: id('1234abcd1234abcd'),
@@ -256,13 +256,14 @@ describe('agentless-ci-visibility-encode', () => {
       },
     ]
 
-    encoder.encode(traceToFilter)
+    encoder.encode(traceWithMixedSpans)
 
     const buffer = encoder.makePayload()
     const decodedTrace = msgpack.decode(buffer, { useBigInt64: true })
-    assert.strictEqual(decodedTrace.events.length, 1)
+    assert.strictEqual(decodedTrace.events.length, 2)
     assert.strictEqual(decodedTrace.events[0].type, 'test_session_end')
     assert.deepStrictEqual(decodedTrace.events[0].content.type, 'test_session_end')
+    assert.strictEqual(decodedTrace.events[1].type, 'span')
   })
 
   it('does not crash if test_session_id is in meta but not test_module_id', () => {
