[test optimization] Fix mocks on jest tests when retried with ATF, EFD or impacted tests (#7352)

juan-fernandez · web-flow · commit 2c57dccc96ae · 2026-01-28T11:01:51.000+01:00
diff --git a/integration-tests/ci-visibility/test-early-flake-detection/test-efd-with-mock.js b/integration-tests/ci-visibility/test-early-flake-detection/test-efd-with-mock.js
@@ -0,0 +1,17 @@
+'use strict'
+
+const mockFn = jest.fn()
+
+describe('early flake detection tests with mock', () => {
+  it('resets mock state between retries', () => {
+    // eslint-disable-next-line no-console
+    console.log('I am running EFD with mock')
+
+    // Call the mock function once
+    mockFn()
+
+    // This assertion should pass on every retry because mock state should be reset
+    // If mock state is NOT reset, this will fail on the 2nd retry with "Expected: 1, Received: 2"
+    expect(mockFn).toHaveBeenCalledTimes(1)
+  })
+})
diff --git a/integration-tests/ci-visibility/test-impacted-test/test-impacted-with-mock.js b/integration-tests/ci-visibility/test-impacted-test/test-impacted-with-mock.js
@@ -0,0 +1,17 @@
+'use strict'
+
+const mockFn = jest.fn()
+
+describe('impacted tests with mock', () => {
+  it('resets mock state between retries', () => {
+    // eslint-disable-next-line no-console
+    console.log('I am running impacted test with mock')
+
+    // Call the mock function once
+    mockFn()
+
+    // This assertion should pass on every retry because mock state should be reset
+    // If mock state is NOT reset, this will fail on the 2nd retry with "Expected: 1, Received: 2"
+    expect(mockFn).toHaveBeenCalledTimes(1)
+  })
+})
diff --git a/integration-tests/ci-visibility/test-management/test-attempt-to-fix-with-mock.js b/integration-tests/ci-visibility/test-management/test-attempt-to-fix-with-mock.js
@@ -0,0 +1,17 @@
+'use strict'
+
+const mockFn = jest.fn()
+
+describe('attempt to fix tests with mock', () => {
+  it('resets mock state between retries', () => {
+    // eslint-disable-next-line no-console
+    console.log('I am running attempt to fix with mock')
+
+    // Call the mock function once
+    mockFn()
+
+    // This assertion should pass on every retry because mock state should be reset
+    // If mock state is NOT reset, this will fail on the 2nd retry with "Expected: 1, Received: 2"
+    expect(mockFn).toHaveBeenCalledTimes(1)
+  })
+})
diff --git a/integration-tests/jest/jest.spec.js b/integration-tests/jest/jest.spec.js
@@ -2170,6 +2170,76 @@ describe(`jest@${JEST_VERSION} commonJS`, () => {
       })
     })
 
+    it('resets mock state between early flake detection retries', async () => {
+      receiver.setInfoResponse({ endpoints: ['/evp_proxy/v4'] })
+      // Test is considered new (not in known tests)
+      receiver.setKnownTests({ jest: {} })
+      const NUM_RETRIES_EFD = 3
+      receiver.setSettings({
+        early_flake_detection: {
+          enabled: true,
+          slow_test_retries: {
+            '5s': NUM_RETRIES_EFD
+          },
+          faulty_session_threshold: 100
+        },
+        known_tests_enabled: true
+      })
+
+      let stdout = ''
+      const eventsPromise = receiver
+        .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/citestcycle'), (payloads) => {
+          const events = payloads.flatMap(({ payload }) => payload.events)
+          const tests = events.filter(event => event.type === 'test').map(event => event.content)
+
+          // Should have 1 original + NUM_RETRIES_EFD retry attempts
+          const mockTests = tests.filter(
+            test => test.meta[TEST_NAME] === 'early flake detection tests with mock resets mock state between retries'
+          )
+          assert.strictEqual(mockTests.length, NUM_RETRIES_EFD + 1)
+
+          // All tests should pass because mock state is reset between retries
+          for (const test of mockTests) {
+            assert.strictEqual(test.meta[TEST_STATUS], 'pass')
+          }
+
+          // All should be marked as new
+          for (const test of mockTests) {
+            assert.strictEqual(test.meta[TEST_IS_NEW], 'true')
+          }
+        })
+
+      childProcess = exec(
+        runTestsCommand,
+        {
+          cwd,
+          env: {
+            ...getCiVisEvpProxyConfig(receiver.port),
+            TESTS_TO_RUN: 'test-early-flake-detection/test-efd-with-mock'
+          },
+        }
+      )
+
+      childProcess.stdout?.on('data', (chunk) => {
+        stdout += chunk.toString()
+      })
+
+      childProcess.stderr?.on('data', (chunk) => {
+        stdout += chunk.toString()
+      })
+
+      const [exitCode] = await Promise.all([
+        once(childProcess, 'exit'),
+        eventsPromise
+      ])
+
+      // Verify the test actually ran
+      assert.match(stdout, /I am running EFD with mock/)
+
+      // All retries should pass, so exit code should be 0
+      assert.strictEqual(exitCode[0], 0)
+    })
+
     it('handles parameterized tests as a single unit', (done) => {
       receiver.setInfoResponse({ endpoints: ['/evp_proxy/v4'] })
       // Tests from ci-visibility/test-early-flake-detection/test-parameterized.js will be considered new
@@ -4208,6 +4278,79 @@ describe(`jest@${JEST_VERSION} commonJS`, () => {
         ])
       })
 
+      it('resets mock state between attempt to fix retries', async () => {
+        const NUM_RETRIES = 3
+        receiver.setSettings({ test_management: { enabled: true, attempt_to_fix_retries: NUM_RETRIES } })
+
+        receiver.setTestManagementTests({
+          jest: {
+            suites: {
+              'ci-visibility/test-management/test-attempt-to-fix-with-mock.js': {
+                tests: {
+                  'attempt to fix tests with mock resets mock state between retries': {
+                    properties: {
+                      attempt_to_fix: true
+                    }
+                  }
+                }
+              }
+            }
+          }
+        })
+
+        let stdout = ''
+        const eventsPromise = receiver
+          .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/citestcycle'), (payloads) => {
+            const events = payloads.flatMap(({ payload }) => payload.events)
+            const tests = events.filter(event => event.type === 'test').map(event => event.content)
+
+            // Should have 1 original + NUM_RETRIES retry attempts
+            const mockTests = tests.filter(
+              test => test.meta[TEST_NAME] === 'attempt to fix tests with mock resets mock state between retries'
+            )
+            assert.strictEqual(mockTests.length, NUM_RETRIES + 1)
+
+            // All tests should pass because mock state is reset between retries
+            for (const test of mockTests) {
+              assert.strictEqual(test.meta[TEST_STATUS], 'pass')
+            }
+
+            // Last attempt should be marked as attempt_to_fix_passed
+            const lastTest = mockTests[mockTests.length - 1]
+            assert.strictEqual(lastTest.meta[TEST_MANAGEMENT_ATTEMPT_TO_FIX_PASSED], 'true')
+          })
+
+        childProcess = exec(
+          runTestsCommand,
+          {
+            cwd,
+            env: {
+              ...getCiVisAgentlessConfig(receiver.port),
+              TESTS_TO_RUN: 'test-management/test-attempt-to-fix-with-mock'
+            },
+          }
+        )
+
+        childProcess.stdout?.on('data', (chunk) => {
+          stdout += chunk.toString()
+        })
+
+        childProcess.stderr?.on('data', (chunk) => {
+          stdout += chunk.toString()
+        })
+
+        const [exitCode] = await Promise.all([
+          once(childProcess, 'exit'),
+          eventsPromise
+        ])
+
+        // Verify the test actually ran
+        assert.match(stdout, /I am running attempt to fix with mock/)
+
+        // All retries should pass, so exit code should be 0
+        assert.strictEqual(exitCode[0], 0)
+      })
+
       it('does not fail retry if a test is quarantined', (done) => {
         receiver.setSettings({ test_management: { enabled: true, attempt_to_fix_retries: 3 } })
         receiver.setTestManagementTests({
@@ -4936,6 +5079,24 @@ describe(`jest@${JEST_VERSION} commonJS`, () => {
          })`
       )
       execSync('git add ci-visibility/test-impacted-test/test-impacted-1.js', { cwd, stdio: 'ignore' })
+
+      // Also modify test file with mock for mock state reset test
+      fs.writeFileSync(
+        path.join(cwd, 'ci-visibility/test-impacted-test/test-impacted-with-mock.js'),
+        `'use strict'
+
+         const mockFn = jest.fn()
+
+         describe('impacted tests with mock', () => {
+           it('resets mock state between retries', () => {
+             console.log('I am running impacted test with mock')
+             mockFn()
+             expect(mockFn).toHaveBeenCalledTimes(1)
+           })
+         })`
+      )
+      execSync('git add ci-visibility/test-impacted-test/test-impacted-with-mock.js', { cwd, stdio: 'ignore' })
+
       execSync('git commit -m "modify test-impacted-1.js"', { cwd, stdio: 'ignore' })
     })
 
@@ -5161,6 +5322,76 @@ describe(`jest@${JEST_VERSION} commonJS`, () => {
         })
         runImpactedTest(done, { isModified: true, isEfd: true, isNew: true })
       })
+
+      it('resets mock state between impacted test retries', async () => {
+        // Test is considered new (not in known tests)
+        receiver.setKnownTests({ jest: {} })
+        receiver.setSettings({
+          impacted_tests_enabled: true,
+          early_flake_detection: {
+            enabled: true,
+            slow_test_retries: {
+              '5s': NUM_RETRIES
+            },
+            faulty_session_threshold: 100
+          },
+          known_tests_enabled: true
+        })
+
+        let stdout = ''
+        const eventsPromise = receiver
+          .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/citestcycle'), (payloads) => {
+            const events = payloads.flatMap(({ payload }) => payload.events)
+            const tests = events.filter(event => event.type === 'test').map(event => event.content)
+
+            // Should have 1 original + NUM_RETRIES retry attempts
+            const mockTests = tests.filter(
+              test => test.meta[TEST_NAME] === 'impacted tests with mock resets mock state between retries'
+            )
+            assert.strictEqual(mockTests.length, NUM_RETRIES + 1)
+
+            // All tests should pass because mock state is reset between retries
+            for (const test of mockTests) {
+              assert.strictEqual(test.meta[TEST_STATUS], 'pass')
+            }
+
+            // All should be marked as modified (impacted)
+            for (const test of mockTests) {
+              assert.strictEqual(test.meta[TEST_IS_MODIFIED], 'true')
+            }
+          })
+
+        childProcess = exec(
+          runTestsCommand,
+          {
+            cwd,
+            env: {
+              ...getCiVisAgentlessConfig(receiver.port),
+              TESTS_TO_RUN: 'test-impacted-test/test-impacted-with-mock',
+              GITHUB_BASE_REF: ''
+            },
+          }
+        )
+
+        childProcess.stdout?.on('data', (chunk) => {
+          stdout += chunk.toString()
+        })
+
+        childProcess.stderr?.on('data', (chunk) => {
+          stdout += chunk.toString()
+        })
+
+        const [exitCode] = await Promise.all([
+          once(childProcess, 'exit'),
+          eventsPromise
+        ])
+
+        // Verify the test actually ran
+        assert.match(stdout, /I am running impacted test with mock/)
+
+        // All retries should pass, so exit code should be 0
+        assert.strictEqual(exitCode[0], 0)
+      })
     })
   })
 
diff --git a/packages/datadog-instrumentations/src/jest.js b/packages/datadog-instrumentations/src/jest.js
@@ -91,6 +91,7 @@ const wrappedWorkers = new WeakSet()
 const testSuiteMockedFiles = new Map()
 const testsToBeRetried = new Set()
 const testSuiteAbsolutePathsWithFastCheck = new Set()
+const testSuiteJestObjects = new Map()
 
 const BREAKPOINT_HIT_GRACE_PERIOD_MS = 200
 const ATR_RETRY_SUPPRESSION_FLAG = '_ddDisableAtrRetry'
@@ -237,6 +238,28 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
       }
     }
 
+    /**
+     * Jest mock state issue during test retries
+     *
+     * Problem:
+     * - Jest tracks mock function calls using internal state (call count, call arguments, etc.)
+     * - When a test is retried, the mock state is not automatically reset
+     * - This causes assertions like `toHaveBeenCalledTimes(1)` to fail because the call count
+     *   accumulates across retries
+     *
+     * The solution is to clear all mocks before each retry attempt.
+     */
+    resetMockState () {
+      try {
+        const jestObject = testSuiteJestObjects.get(this.testSuiteAbsolutePath)
+        if (jestObject?.clearAllMocks) {
+          jestObject.clearAllMocks()
+        }
+      } catch (e) {
+        log.warn('Error resetting mock state', e)
+      }
+    }
+
     // This function returns an array if the known tests are valid and null otherwise.
     getKnownTestsForSuite (suiteKnownTests) {
       // `suiteKnownTests` is `this.testEnvironmentOptions._ddKnownTests`,
@@ -339,8 +362,9 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
       if (event.name === 'test_start') {
         const testName = getJestTestName(event.test, this.getShouldStripSeedFromTestName())
         if (testsToBeRetried.has(testName)) {
-          // This is needed because we're trying tests with the same name
+          // This is needed because we're retrying tests with the same name
           this.resetSnapshotState()
+          this.resetMockState()
         }
 
         let isNewTest = false
@@ -1148,9 +1172,19 @@ function jestAdapterWrapper (jestAdapter, jestVersion) {
         })
       }
       testSuiteFinishCh.publish({ status, errorMessage, testSuiteAbsolutePath: environment.testSuiteAbsolutePath })
+
+      // Cleanup per-suite state to avoid memory leaks
+      testSuiteMockedFiles.delete(environment.testSuiteAbsolutePath)
+      testSuiteJestObjects.delete(environment.testSuiteAbsolutePath)
+
       return suiteResults
     }).catch(error => {
       testSuiteFinishCh.publish({ status: 'fail', error, testSuiteAbsolutePath: environment.testSuiteAbsolutePath })
+
+      // Cleanup per-suite state to avoid memory leaks
+      testSuiteMockedFiles.delete(environment.testSuiteAbsolutePath)
+      testSuiteJestObjects.delete(environment.testSuiteAbsolutePath)
+
       throw error
     })
   })
@@ -1312,6 +1346,11 @@ addHook({
     const result = _createJestObjectFor.apply(this, arguments)
     const suiteFilePath = this._testPath || from
 
+    // Store the jest object so we can access it later for resetting mock state
+    if (suiteFilePath) {
+      testSuiteJestObjects.set(suiteFilePath, result)
+    }
+
     shimmer.wrap(result, 'mock', mock => function (moduleName) {
       // If the library is mocked with `jest.mock`, we don't want to bypass jest's own require engine
       if (LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE.has(moduleName)) {
