fix(jest): properly parent beforeAll/afterAll hooks to suite span

integration-tests/ci-visibility/jest-hooks/jest-hooks-test.js

@@ -0,0 +1,92 @@
+'use strict'
+
+describe('Jest Hook Instrumentation Test Suite', () => {
+  let suiteData = []
+  let testData = []
+
+  // Suite-level hooks - should be parented to suite span
+  beforeAll(() => {
+    console.log('beforeAll hook executed')
+    suiteData.push('beforeAll')
+  })
+
+  afterAll(() => {
+    console.log('afterAll hook executed')
+    suiteData.push('afterAll')
+  })
+
+  // Test-level hooks - should be parented to test span
+  beforeEach(() => {
+    console.log('beforeEach hook executed')
+    testData.push('beforeEach')
+  })
+
+  afterEach(() => {
+    console.log('afterEach hook executed')
+    testData.push('afterEach')
+  })
+
+  test('test with all hooks', () => {
+    console.log('test executed')
+    expect(true).toBe(true)
+  })
+
+  test('second test to verify hook execution', () => {
+    console.log('second test executed')
+    expect(testData.length).toBeGreaterThan(0)
+  })
+
+  describe('nested describe block', () => {
+    beforeAll(() => {
+      console.log('nested beforeAll hook executed')
+      suiteData.push('nested-beforeAll')
+    })
+
+    afterAll(() => {
+      console.log('nested afterAll hook executed')
+      suiteData.push('nested-afterAll')
+    })
+
+    test('nested test', () => {
+      console.log('nested test executed')
+      expect(suiteData.length).toBeGreaterThan(0)
+    })
+  })
+})
+
+describe('Async Hook Test Suite', () => {
+  let asyncData = null
+
+  beforeAll(async () => {
+    console.log('async beforeAll hook started')
+    await new Promise(resolve => setTimeout(resolve, 10))
+    asyncData = 'initialized'
+    console.log('async beforeAll hook completed')
+  })
+
+  afterAll(async () => {
+    console.log('async afterAll hook started')
+    await new Promise(resolve => setTimeout(resolve, 10))
+    asyncData = null
+    console.log('async afterAll hook completed')
+  })
+
+  test('test with async hooks', () => {
+    console.log('test with async data:', asyncData)
+    expect(asyncData).toBe('initialized')
+  })
+})
+
+describe('Hook Error Test Suite', () => {
+  // This hook should fail and be attributed to the suite
+  beforeAll(() => {
+    console.log('beforeAll hook that will fail')
+    // Uncomment to test error handling
+    // throw new Error('beforeAll hook error')
+  })
+
+  test('test that should run if beforeAll succeeds', () => {
+    console.log('test after beforeAll')
+    expect(true).toBe(true)
+  })
+})

integration-tests/ci-visibility/jest-hooks/run-jest-hooks.js

@@ -0,0 +1,27 @@
+#!/usr/bin/env node
+'use strict'
+
+const jest = require('jest')
+
+const options = {
+  projects: [__dirname],
+  testMatch: ['**/jest-hooks-test.js'],
+  coverageReporters: ['json'],
+  coverage: false,
+  maxWorkers: 1,
+  testEnvironment: 'node'
+}
+
+jest.runCLI(options, options.projects)
+  .then((result) => {
+    if (result.results.success) {
+      console.log('All tests passed!')
+    } else {
+      console.log('Some tests failed')
+      process.exit(1)
+    }
+  })
+  .catch((error) => {
+    console.error('Error running tests:', error)
+    process.exit(1)
+  })

integration-tests/jest/jest.spec.js

@@ -5093,4 +5093,72 @@ describe(`jest@${JEST_VERSION} commonJS`, () => {
     assert.notInclude(testOutput, 'Cannot find module')
     assert.include(testOutput, '6 passed')
   })
+
+  context('hook instrumentation', () => {
+    it('should properly parent beforeAll and afterAll hooks to suite span', async () => {
+      const hookTestPromise = receiver.gatherPayloadsMaxTimeout(({ url }) => url === '/api/v2/citestcycle', 5000)
+
+      childProcess = exec(
+        'node ./ci-visibility/jest-hooks/run-jest-hooks.js',
+        {
+          cwd,
+          env: {
+            ...getCiVisAgentlessConfig(receiver.port),
+            NODE_OPTIONS: '-r dd-trace/ci/init'
+          },
+          stdio: 'pipe'
+        }
+      )
+
+      const [code] = await once(childProcess, 'exit')
+      assert.equal(code, 0, 'Jest hook tests should pass')
+
+      const payloads = await hookTestPromise
+      const events = payloads.flatMap(({ payload }) => payload.events)
+
+      // Find test suite events
+      const suiteEvents = events.filter(event => event.type === 'test_suite')
+      assert.isAbove(suiteEvents.length, 0, 'Should have test suite events')
+
+      // Find test events
+      const testEvents = events.filter(event => event.type === 'test')
+      assert.isAbove(testEvents.length, 0, 'Should have test events')
+
+      // Verify that hooks executed successfully
+      const passedTests = testEvents.filter(test => test.content.meta[TEST_STATUS] === 'pass')
+      assert.isAbove(passedTests.length, 0, 'Should have passing tests')
+
+      // Check that all test suites passed
+      const passedSuites = suiteEvents.filter(suite => suite.content.meta[TEST_STATUS] === 'pass')
+      assert.equal(passedSuites.length, suiteEvents.length, 'All test suites should pass')
+    })
+
+    it('should handle async hooks correctly', async () => {
+      const asyncHookTestPromise = receiver.gatherPayloadsMaxTimeout(({ url }) => url === '/api/v2/citestcycle', 5000)
+
+      childProcess = exec(
+        'node ./ci-visibility/jest-hooks/run-jest-hooks.js',
+        {
+          cwd,
+          env: {
+            ...getCiVisAgentlessConfig(receiver.port),
+            NODE_OPTIONS: '-r dd-trace/ci/init',
+            TESTS_TO_RUN: 'Async Hook Test Suite'
+          },
+          stdio: 'pipe'
+        }
+      )
+
+      const [code] = await once(childProcess, 'exit')
+      assert.equal(code, 0, 'Async hook tests should pass')
+
+      const payloads = await asyncHookTestPromise
+      const events = payloads.flatMap(({ payload }) => payload.events)
+      const testEvents = events.filter(event => event.type === 'test')
+      const asyncTests = testEvents.filter(test => test.content.meta[TEST_NAME]?.includes('async hooks'))
+
+      assert.isAbove(asyncTests.length, 0, 'Should have async hook tests')
+      assert.equal(asyncTests.every(test => test.content.meta[TEST_STATUS] === 'pass'), true, 'Async hook tests should pass')
+    })
+  })
 })

packages/datadog-instrumentations/src/jest.js

@@ -43,6 +43,9 @@ const testFinishCh = channel('ci:jest:test:finish')
 const testErrCh = channel('ci:jest:test:err')
 const testFnCh = channel('ci:jest:test:fn')
 
+// Suite-level channels for hook execution
+const suiteFnCh = channel('ci:jest:suite:fn')
+
 const skippableSuitesCh = channel('ci:jest:test-suite:skippable')
 const libraryConfigurationCh = channel('ci:jest:library-configuration')
 const knownTestsCh = channel('ci:jest:known-tests')
@@ -84,6 +87,7 @@ let modifiedFiles = {}
 const testContexts = new WeakMap()
 const originalTestFns = new WeakMap()
 const originalHookFns = new WeakMap()
+const suiteContexts = new Map() // Map from test suite path to suite context
 const retriedTestsToNumAttempts = new Map()
 const newTestsTestStatuses = new Map()
 const attemptToFixRetriedTestsStatuses = new Map()
@@ -121,6 +125,18 @@ function getTestEnvironmentOptions (config) {
   return {}
 }
 
+// Helper functions to identify hook types
+function getHookType (hook) {
+  // Jest Circus stores hook type in hook.type
+  // Values: 'beforeAll', 'afterAll', 'beforeEach', 'afterEach'
+  return hook.type
+}
+
+function isSuiteLevelHook (hook) {
+  const type = getHookType(hook)
+  return type === 'beforeAll' || type === 'afterAll'
+}
+
 function getTestStats (testStatuses) {
   return testStatuses.reduce((acc, testStatus) => {
     acc[testStatus]++
@@ -139,6 +155,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
       this.global._ddtrace = global._ddtrace
       this.hasSnapshotTests = undefined
       this.testSuiteAbsolutePath = context.testPath
+      this.suiteSpanContext = null // Will store suite span context for hook wrapping
 
       this.displayName = config.projectConfig?.displayName?.name || config.displayName
       this.testEnvironmentOptions = getTestEnvironmentOptions(config)
@@ -323,6 +340,15 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
 
       const setNameToParams = (name, params) => { this.nameToParams[name] = [...params] }
 
+      // Capture suite context when a describe block starts
+      if (event.name === 'run_describe_start') {
+        // Store the suite context for this test suite
+        const suiteContext = suiteContexts.get(this.testSuiteAbsolutePath)
+        if (suiteContext) {
+          this.suiteSpanContext = suiteContext
+        }
+      }
+
       if (event.name === 'setup' && this.global.test) {
         shimmer.wrap(this.global.test, 'each', each => function () {
           const testParameters = getFormattedJestTestParameters(arguments)
@@ -405,14 +431,27 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
 
         testStartCh.runStores(ctx, () => {
           for (const hook of event.test.parent.hooks) {
-            let hookFn = hook.fn
+            // Skip if hook is already wrapped
             if (originalHookFns.has(hook)) {
-              hookFn = originalHookFns.get(hook)
-            } else {
-              originalHookFns.set(hook, hookFn)
+              continue
             }
-            const newHookFn = shimmer.wrapFunction(hookFn, hookFn => function () {
-              return testFnCh.runStores(ctx, () => hookFn.apply(this, arguments))
+
+            const hookType = getHookType(hook)
+            const isSuiteHook = isSuiteLevelHook(hook)
+
+            // Store original function
+            originalHookFns.set(hook, hook.fn)
+
+            // Choose appropriate context based on hook type
+            const hookContext = isSuiteHook ? this.suiteSpanContext : ctx
+
+            // Wrap with appropriate context
+            const newHookFn = shimmer.wrapFunction(hook.fn, hookFn => function () {
+              // Use different channel for suite hooks vs test hooks
+              const channel = isSuiteHook ? suiteFnCh : testFnCh
+              // If suite context is not available, fall back to test context
+              const finalContext = hookContext || ctx
+              return channel.runStores(finalContext, () => hookFn.apply(this, arguments))
             })
             hook.fn = newHookFn
           }
@@ -1089,13 +1128,22 @@ function jestAdapterWrapper (jestAdapter, jestVersion) {
     if (!environment || !environment.testEnvironmentOptions) {
       return adapter.apply(this, arguments)
     }
-    testSuiteStartCh.publish({
+    const suiteContext = {
       testSuite: environment.testSuite,
       testEnvironmentOptions: environment.testEnvironmentOptions,
       testSourceFile: environment.testSourceFile,
       displayName: environment.displayName,
       frameworkVersion: jestVersion,
       testSuiteAbsolutePath: environment.testSuiteAbsolutePath
+    }
+
+    // Publish suite start to create the span, then capture its context
+    testSuiteStartCh.runStores(suiteContext, () => {
+      // Store suite context with the span information for hook wrapping
+      const store = require('../../datadog-core').storage('legacy').getStore()
+      if (store) {
+        suiteContexts.set(environment.testSuiteAbsolutePath, store)
+      }
     })
     return adapter.apply(this, arguments).then(suiteResults => {
       const { numFailingTests, skipped, failureMessage: errorMessage } = suiteResults

packages/datadog-plugin-jest/src/index.js

@@ -193,14 +193,15 @@ class JestPlugin extends CiPlugin {
       })
     })
 
-    this.addSub('ci:jest:test-suite:start', ({
+    this.addBind('ci:jest:test-suite:start', (ctx) => {
+      const {
       testSuite,
       testSourceFile,
       testEnvironmentOptions,
       frameworkVersion,
       displayName,
       testSuiteAbsolutePath
-    }) => {
+    } = ctx
       const {
         _ddTestSessionId: testSessionId,
         _ddTestCommand: testCommand,
@@ -263,6 +264,10 @@ class JestPlugin extends CiPlugin {
         this.telemetry.ciVisEvent(TELEMETRY_CODE_COVERAGE_STARTED, 'suite', { library: 'istanbul' })
       }
       this.testSuiteSpanPerTestSuiteAbsolutePath.set(testSuiteAbsolutePath, this.testSuiteSpan)
+
+      // Return the store with the suite span for context
+      const store = storage('legacy').getStore()
+      return { ...store, span: this.testSuiteSpan }
     })
 
     this.addSub('ci:jest:worker-report:coverage', data => {
@@ -360,6 +365,12 @@ class JestPlugin extends CiPlugin {
       return ctx.currentStore
     })
 
+    // Add binding for suite function channel (used for suite-level hooks)
+    this.addBind('ci:jest:suite:fn', (ctx) => {
+      // Return suite context for hook execution
+      return ctx.currentStore || ctx
+    })
+
     this.addSub('ci:jest:test:finish', ({
       span,
       status,