[test optimization] Fix attempt to fix in vitest (#7785)

juan-fernandez · web-flow · commit 453f900bf1ac · 2026-03-16T13:26:58.000+01:00
diff --git a/integration-tests/ci-visibility/vitest-tests/test-attempt-to-fix.mjs b/integration-tests/ci-visibility/vitest-tests/test-attempt-to-fix.mjs
@@ -15,6 +15,14 @@ describe('attempt to fix tests', () => {
       } else {
         expect(1 + 2).to.equal(3)
       }
+    } else if (process.env.SHOULD_FAIL_FIRST_ONLY) {
+      // First attempt fails, all retries pass. Exit code must still be 1
+      // for plain ATF tests (not quarantined/disabled).
+      if (numAttempt++ === 0) {
+        expect(1 + 2).to.equal(4)
+      } else {
+        expect(1 + 2).to.equal(3)
+      }
     } else {
       expect(1 + 2).to.equal(4)
     }
diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
@@ -1716,6 +1716,7 @@ versions.forEach((version) => {
             shouldAlwaysPass,
             isQuarantining,
             shouldFailSometimes,
+            shouldFailFirstOnly,
             isDisabling,
             extraEnvVars = {},
           } = {}) => {
@@ -1738,6 +1739,7 @@ versions.forEach((version) => {
                   ...extraEnvVars,
                   ...(shouldAlwaysPass ? { SHOULD_ALWAYS_PASS: '1' } : {}),
                   ...(shouldFailSometimes ? { SHOULD_FAIL_SOMETIMES: '1' } : {}),
+                  ...(shouldFailFirstOnly ? { SHOULD_FAIL_FIRST_ONLY: '1' } : {}),
                 },
               }
             )
@@ -1777,6 +1779,12 @@ versions.forEach((version) => {
             runAttemptToFixTest(done, { isAttemptingToFix: true, shouldFailSometimes: true })
           })
 
+          it('does not suppress exit code for plain ATF tests even when last retry passes', (done) => {
+            receiver.setSettings({ test_management: { enabled: true, attempt_to_fix_retries: 3 } })
+
+            runAttemptToFixTest(done, { isAttemptingToFix: true, shouldFailFirstOnly: true })
+          })
+
           it('does not attempt to fix tests if test management is not enabled', (done) => {
             receiver.setSettings({ test_management: { enabled: false, attempt_to_fix_retries: 3 } })
 
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
@@ -747,7 +747,10 @@ function wrapVitestTestRunner (VitestTestRunner) {
     }
 
     const lastExecutionStatus = task.result.state
-    const shouldFlipStatus = isEarlyFlakeDetectionEnabled || attemptToFixTasks.has(task)
+    const isAtf = attemptToFixTasks.has(task)
+    const isQuarantinedOrDisabledAtf = isAtf && (quarantinedTasks.has(task) || disabledTasks.has(task))
+    const shouldTrackStatuses = isEarlyFlakeDetectionEnabled || isAtf
+    const shouldFlipStatus = isEarlyFlakeDetectionEnabled || isQuarantinedOrDisabledAtf
     const statuses = taskToStatuses.get(task)
 
     // These clauses handle task.repeats, whether EFD is enabled or not
@@ -765,8 +768,10 @@ function wrapVitestTestRunner (VitestTestRunner) {
         } else {
           testPassCh.publish({ task, ...ctx.currentStore })
         }
-        if (shouldFlipStatus) {
+        if (shouldTrackStatuses) {
           statuses.push(lastExecutionStatus)
+        }
+        if (shouldFlipStatus) {
           // If we don't "reset" the result.state to "pass", once a repetition fails,
           // vitest will always consider the test as failed, so we can't read the actual status
           // This means that we change vitest's behavior:
@@ -776,7 +781,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
         }
       }
     } else if (numRepetition === task.repeats) {
-      if (shouldFlipStatus) {
+      if (shouldTrackStatuses) {
         statuses.push(lastExecutionStatus)
       }
 

Original file line number	Diff line number	Diff line change
`@@ -747,7 +747,10 @@ function wrapVitestTestRunner (VitestTestRunner) {`
`747`	`747`	`}`
`748`	`748`
`749`	`749`	`const lastExecutionStatus = task.result.state`
`750`		`- const shouldFlipStatus = isEarlyFlakeDetectionEnabled \|\| attemptToFixTasks.has(task)`
	`750`	`+ const isAtf = attemptToFixTasks.has(task)`
	`751`	`+ const isQuarantinedOrDisabledAtf = isAtf && (quarantinedTasks.has(task) \|\| disabledTasks.has(task))`
	`752`	`+ const shouldTrackStatuses = isEarlyFlakeDetectionEnabled \|\| isAtf`
	`753`	`+ const shouldFlipStatus = isEarlyFlakeDetectionEnabled \|\| isQuarantinedOrDisabledAtf`
`751`	`754`	`const statuses = taskToStatuses.get(task)`
`752`	`755`
`753`	`756`	`// These clauses handle task.repeats, whether EFD is enabled or not`
`@@ -765,8 +768,10 @@ function wrapVitestTestRunner (VitestTestRunner) {`
`765`	`768`	`} else {`
`766`	`769`	`testPassCh.publish({ task, ...ctx.currentStore })`
`767`	`770`	`}`
`768`		`- if (shouldFlipStatus) {`
	`771`	`+ if (shouldTrackStatuses) {`
`769`	`772`	`statuses.push(lastExecutionStatus)`
	`773`	`+ }`
	`774`	`+ if (shouldFlipStatus) {`
`770`	`775`	`// If we don't "reset" the result.state to "pass", once a repetition fails,`
`771`	`776`	`// vitest will always consider the test as failed, so we can't read the actual status`
`772`	`777`	`// This means that we change vitest's behavior:`
`@@ -776,7 +781,7 @@ function wrapVitestTestRunner (VitestTestRunner) {`
`776`	`781`	`}`
`777`	`782`	`}`
`778`	`783`	`} else if (numRepetition === task.repeats) {`
`779`		`- if (shouldFlipStatus) {`
	`784`	`+ if (shouldTrackStatuses) {`
`780`	`785`	`statuses.push(lastExecutionStatus)`
`781`	`786`	`}`
`782`	`787`