dd-trace/packages/datadog-instrumentations/src/vitest.js

Datadog APM tracing client for JavaScript

'use strict'

const shimmer = require('../../datadog-shimmer')
const log = require('../../dd-trace/src/log')
const {
  VITEST_WORKER_TRACE_PAYLOAD_CODE,
  VITEST_WORKER_LOGS_PAYLOAD_CODE
} = require('../../dd-trace/src/plugins/util/test')
const { addHook, channel } = require('./helpers/instrument')

// test hooks
const testStartCh = channel('ci:vitest:test:start')
const testFinishTimeCh = channel('ci:vitest:test:finish-time')
const testPassCh = channel('ci:vitest:test:pass')
const testErrorCh = channel('ci:vitest:test:error')
const testSkipCh = channel('ci:vitest:test:skip')
const isNewTestCh = channel('ci:vitest:test:is-new')
const isAttemptToFixCh = channel('ci:vitest:test:is-attempt-to-fix')
const isDisabledCh = channel('ci:vitest:test:is-disabled')
const isQuarantinedCh = channel('ci:vitest:test:is-quarantined')
const isModifiedCh = channel('ci:vitest:test:is-modified')

// test suite hooks
const testSuiteStartCh = channel('ci:vitest:test-suite:start')
const testSuiteFinishCh = channel('ci:vitest:test-suite:finish')
const testSuiteErrorCh = channel('ci:vitest:test-suite:error')

// test session hooks
const testSessionStartCh = channel('ci:vitest:session:start')
const testSessionFinishCh = channel('ci:vitest:session:finish')
const libraryConfigurationCh = channel('ci:vitest:library-configuration')
const knownTestsCh = channel('ci:vitest:known-tests')
const isEarlyFlakeDetectionFaultyCh = channel('ci:vitest:is-early-flake-detection-faulty')
const testManagementTestsCh = channel('ci:vitest:test-management-tests')
const modifiedFilesCh = channel('ci:vitest:modified-files')

const workerReportTraceCh = channel('ci:vitest:worker-report:trace')
const workerReportLogsCh = channel('ci:vitest:worker-report:logs')

const taskToCtx = new WeakMap()
const taskToStatuses = new WeakMap()
const newTasks = new WeakSet()
const disabledTasks = new WeakSet()
const quarantinedTasks = new WeakSet()
const attemptToFixTasks = new WeakSet()
const modifiedTasks = new WeakSet()
let isRetryReasonEfd = false
let isRetryReasonAttemptToFix = false
const switchedStatuses = new WeakSet()
const workerProcesses = new WeakSet()
let isFlakyTestRetriesEnabled = false
let flakyTestRetriesCount = 0
let isEarlyFlakeDetectionEnabled = false
let earlyFlakeDetectionNumRetries = 0
let isEarlyFlakeDetectionFaulty = false
let isKnownTestsEnabled = false
let isTestManagementTestsEnabled = false
let isImpactedTestsEnabled = false
let testManagementAttemptToFixRetries = 0
let isDiEnabled = false
let testCodeCoverageLinesTotal
let isSessionStarted = false
let vitestPool = null

const BREAKPOINT_HIT_GRACE_PERIOD_MS = 400

function getTestCommand () {
  return `vitest ${process.argv.slice(2).join(' ')}`
}

function waitForHitProbe () {
  return new Promise(resolve => {
    setTimeout(() => {
      resolve()
    }, BREAKPOINT_HIT_GRACE_PERIOD_MS)
  })
}

function isValidKnownTests (receivedKnownTests) {
  return !!receivedKnownTests.vitest
}

function getProvidedContext () {
  try {
    const {
      _ddIsEarlyFlakeDetectionEnabled,
      _ddIsDiEnabled,
      _ddKnownTests: knownTests,
      _ddEarlyFlakeDetectionNumRetries: numRepeats,
      _ddIsKnownTestsEnabled: isKnownTestsEnabled,
      _ddIsTestManagementTestsEnabled: isTestManagementTestsEnabled,
      _ddTestManagementAttemptToFixRetries: testManagementAttemptToFixRetries,
      _ddTestManagementTests: testManagementTests,
      _ddIsFlakyTestRetriesEnabled: isFlakyTestRetriesEnabled,
      _ddIsImpactedTestsEnabled: isImpactedTestsEnabled,
      _ddModifiedFiles: modifiedFiles
    } = globalThis.__vitest_worker__.providedContext

    return {
      isDiEnabled: _ddIsDiEnabled,
      isEarlyFlakeDetectionEnabled: _ddIsEarlyFlakeDetectionEnabled,
      knownTests,
      numRepeats,
      isKnownTestsEnabled,
      isTestManagementTestsEnabled,
      testManagementAttemptToFixRetries,
      testManagementTests,
      isFlakyTestRetriesEnabled,
      isImpactedTestsEnabled,
      modifiedFiles
    }
  } catch {
    log.error('Vitest workers could not parse provided context, so some features will not work.')
    return {
      isDiEnabled: false,
      isEarlyFlakeDetectionEnabled: false,
      knownTests: {},
      numRepeats: 0,
      isKnownTestsEnabled: false,
      isTestManagementTestsEnabled: false,
      testManagementAttemptToFixRetries: 0,
      testManagementTests: {},
      isFlakyTestRetriesEnabled: false,
      isImpactedTestsEnabled: false,
      modifiedFiles: {}
    }
  }
}

function isReporterPackage (vitestPackage) {
  return vitestPackage.B?.name === 'BaseSequencer'
}

// from 2.0.0
function isReporterPackageNew (vitestPackage) {
  return vitestPackage.e?.name === 'BaseSequencer'
}

function isReporterPackageNewest (vitestPackage) {
  return vitestPackage.h?.name === 'BaseSequencer'
}

function isBaseSequencer (vitestPackage) {
  return vitestPackage.b?.name === 'BaseSequencer'
}

function getChannelPromise (channelToPublishTo, frameworkVersion) {
  return new Promise(resolve => {
    channelToPublishTo.publish({ onDone: resolve, frameworkVersion })
  })
}

function isCliApiPackage (vitestPackage) {
  return vitestPackage.s?.name === 'startVitest'
}

function isTestPackage (testPackage) {
  return testPackage.V?.name === 'VitestTestRunner'
}

function hasForksPoolWorker (vitestPackage) {
  return vitestPackage.f?.name === 'ForksPoolWorker'
}

function hasThreadsPoolWorker (vitestPackage) {
  return vitestPackage.T?.name === 'ThreadsPoolWorker'
}

function getSessionStatus (state) {
  if (state.getCountOfFailedTests() > 0) {
    return 'fail'
  }
  if (state.pathsSet.size === 0) {
    return 'skip'
  }
  return 'pass'
}

// From https://github.com/vitest-dev/vitest/blob/51c04e2f44d91322b334f8ccbcdb368facc3f8ec/packages/runner/src/run.ts#L243-L250
function getVitestTestStatus (test, retryCount) {
  if (test.result.state !== 'fail') {
    if (!test.repeats) {
      return 'pass'
    } else if (test.repeats && (test.retry ?? 0) === retryCount) {
      return 'pass'
    }
  }
  return 'fail'
}

function getTypeTasks (fileTasks, type = 'test') {
  const typeTasks = []

  function getTasks (tasks) {
    for (const task of tasks) {
      if (task.type === type) {
        typeTasks.push(task)
      } else if (task.tasks) {
        getTasks(task.tasks)
      }
    }
  }

  getTasks(fileTasks)

  return typeTasks
}

function getTestName (task) {
  let testName = task.name
  let currentTask = task.suite

  while (currentTask) {
    if (currentTask.name) {
      testName = `${currentTask.name} ${testName}`
    }
    currentTask = currentTask.suite
  }

  return testName
}

function getSortWrapper (sort, frameworkVersion) {
  return async function () {
    if (!testSessionFinishCh.hasSubscribers) {
      return sort.apply(this, arguments)
    }
    // There isn't any other async function that we seem to be able to hook into
    // So we will use the sort from BaseSequencer. This means that a custom sequencer
    // will not work. This will be a known limitation.

    try {
      const { err, libraryConfig } = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
      if (!err) {
        isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
        flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
        isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
        earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
        isDiEnabled = libraryConfig.isDiEnabled
        isKnownTestsEnabled = libraryConfig.isKnownTestsEnabled
        isTestManagementTestsEnabled = libraryConfig.isTestManagementEnabled
        testManagementAttemptToFixRetries = libraryConfig.testManagementAttemptToFixRetries
        isImpactedTestsEnabled = libraryConfig.isImpactedTestsEnabled
      }
    } catch {
      isFlakyTestRetriesEnabled = false
      isEarlyFlakeDetectionEnabled = false
      isDiEnabled = false
      isKnownTestsEnabled = false
      isImpactedTestsEnabled = false
    }

    if (isFlakyTestRetriesEnabled && !this.ctx.config.retry && flakyTestRetriesCount > 0) {
      this.ctx.config.retry = flakyTestRetriesCount
      try {
        const workspaceProject = this.ctx.getCoreWorkspaceProject
          ? this.ctx.getCoreWorkspaceProject()
          : this.ctx.getRootProject()
        workspaceProject._provided._ddIsFlakyTestRetriesEnabled = isFlakyTestRetriesEnabled
      } catch {
        log.warn('Could not send library configuration to workers.')
      }
    }

    if (isKnownTestsEnabled) {
      const knownTestsResponse = await getChannelPromise(knownTestsCh)
      if (knownTestsResponse.err) {
        isEarlyFlakeDetectionEnabled = false
      } else {
        const knownTests = knownTestsResponse.knownTests
        const getFilePaths = this.ctx.getTestFilepaths || this.ctx._globTestFilepaths

        const testFilepaths = await getFilePaths.call(this.ctx)

        if (isValidKnownTests(knownTests)) {
          isEarlyFlakeDetectionFaultyCh.publish({
            knownTests: knownTests.vitest,
            testFilepaths,
            onDone: (isFaulty) => {
              isEarlyFlakeDetectionFaulty = isFaulty
            }
          })
          if (isEarlyFlakeDetectionFaulty) {
            isEarlyFlakeDetectionEnabled = false
            log.warn('New test detection is disabled because the number of new tests is too high.')
          } else {
            // TODO: use this to pass session and module IDs to the worker, instead of polluting process.env
            // Note: setting this.ctx.config.provide directly does not work because it's cached
            try {
              const workspaceProject = this.ctx.getCoreWorkspaceProject
                ? this.ctx.getCoreWorkspaceProject()
                : this.ctx.getRootProject()
              workspaceProject._provided._ddIsKnownTestsEnabled = isKnownTestsEnabled
              workspaceProject._provided._ddKnownTests = knownTests
              workspaceProject._provided._ddIsEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
              workspaceProject._provided._ddEarlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
            } catch {
              log.warn('Could not send known tests to workers so Early Flake Detection will not work.')
            }
          }
        } else {
          isEarlyFlakeDetectionFaulty = true
          isEarlyFlakeDetectionEnabled = false
        }
      }
    }

    if (isDiEnabled) {
      try {
        const workspaceProject = this.ctx.getCoreWorkspaceProject
          ? this.ctx.getCoreWorkspaceProject()
          : this.ctx.getRootProject()
        workspaceProject._provided._ddIsDiEnabled = isDiEnabled
      } catch {
        log.warn('Could not send Dynamic Instrumentation configuration to workers.')
      }
    }

    if (isTestManagementTestsEnabled) {
      const { err, testManagementTests: receivedTestManagementTests } = await getChannelPromise(testManagementTestsCh)
      if (err) {
        isTestManagementTestsEnabled = false
        log.error('Could not get test management tests.')
      } else {
        const testManagementTests = receivedTestManagementTests
        try {
          const workspaceProject = this.ctx.getCoreWorkspaceProject
            ? this.ctx.getCoreWorkspaceProject()
            : this.ctx.getRootProject()
          workspaceProject._provided._ddIsTestManagementTestsEnabled = isTestManagementTestsEnabled
          workspaceProject._provided._ddTestManagementAttemptToFixRetries = testManagementAttemptToFixRetries
          workspaceProject._provided._ddTestManagementTests = testManagementTests
        } catch {
          log.warn('Could not send test management tests to workers so Test Management will not work.')
        }
      }
    }

    if (isImpactedTestsEnabled) {
      const { err, modifiedFiles } = await getChannelPromise(modifiedFilesCh)
      if (err) {
        log.error('Could not get modified tests.')
      } else {
        try {
          const workspaceProject = this.ctx.getCoreWorkspaceProject
            ? this.ctx.getCoreWorkspaceProject()
            : this.ctx.getRootProject()
          workspaceProject._provided._ddIsImpactedTestsEnabled = isImpactedTestsEnabled
          workspaceProject._provided._ddModifiedFiles = modifiedFiles
        } catch {
          log.warn('Could not send modified tests to workers so Impacted Tests will not work.')
        }
      }
    }

    if (this.ctx.coverageProvider?.generateCoverage) {
      shimmer.wrap(this.ctx.coverageProvider, 'generateCoverage', generateCoverage => async function () {
        const totalCodeCoverage = await generateCoverage.apply(this, arguments)

        try {
          testCodeCoverageLinesTotal = totalCodeCoverage.getCoverageSummary().lines.pct
        } catch {
          // ignore errors
        }
        return totalCodeCoverage
      })
    }

    shimmer.wrap(this.ctx, 'exit', getFinishWrapper)
    shimmer.wrap(this.ctx, 'close', getFinishWrapper)

    return sort.apply(this, arguments)
  }
}

function getFinishWrapper (exitOrClose) {
  let isClosed = false
  return async function () {
    if (isClosed) { // needed because exit calls close
      return exitOrClose.apply(this, arguments)
    }
    isClosed = true
    let onFinish

    const flushPromise = new Promise(resolve => {
      onFinish = resolve
    })
    const failedSuites = this.state.getFailedFilepaths()
    let error
    if (failedSuites.length) {
      error = new Error(`Test suites failed: ${failedSuites.length}.`)
    }

    testSessionFinishCh.publish({
      status: getSessionStatus(this.state),
      testCodeCoverageLinesTotal,
      error,
      isEarlyFlakeDetectionEnabled,
      isEarlyFlakeDetectionFaulty,
      isTestManagementTestsEnabled,
      vitestPool,
      onFinish
    })

    await flushPromise

    return exitOrClose.apply(this, arguments)
  }
}

function getCliOrStartVitestWrapper (frameworkVersion) {
  return function (oldCliOrStartVitest) {
    return function () {
      if (!testSessionStartCh.hasSubscribers || isSessionStarted) {
        return oldCliOrStartVitest.apply(this, arguments)
      }
      isSessionStarted = true
      testSessionStartCh.publish({ command: getTestCommand(), frameworkVersion })
      return oldCliOrStartVitest.apply(this, arguments)
    }
  }
}

function getCreateCliWrapper (vitestPackage, frameworkVersion) {
  shimmer.wrap(vitestPackage, 'c', getCliOrStartVitestWrapper(frameworkVersion))

  return vitestPackage
}

function threadHandler (thread) {
  const { runtime } = thread
  let workerProcess
  if (runtime === 'child_process') {
    vitestPool = 'child_process'
    workerProcess = thread.process
  } else if (runtime === 'worker_threads') {
    vitestPool = 'worker_threads'
    workerProcess = thread.thread
  } else {
    vitestPool = 'unknown'
  }
  if (!workerProcess) {
    log.error('Vitest error: could not get process or thread from TinyPool#run')
    return
  }

  if (workerProcesses.has(workerProcess)) {
    return
  }
  workerProcesses.add(workerProcess)
  workerProcess.on('message', (message) => {
    if (message.__tinypool_worker_message__ && message.data) {
      if (message.interprocessCode === VITEST_WORKER_TRACE_PAYLOAD_CODE) {
        workerReportTraceCh.publish(message.data)
      } else if (message.interprocessCode === VITEST_WORKER_LOGS_PAYLOAD_CODE) {
        workerReportLogsCh.publish(message.data)
      }
    }
  })
}

function wrapTinyPoolRun (TinyPool) {
  shimmer.wrap(TinyPool.prototype, 'run', run => async function () {
    // We have to do this before and after because the threads list gets recycled, that is, the processes are re-created
    this.threads.forEach(threadHandler)
    const runResult = await run.apply(this, arguments)
    this.threads.forEach(threadHandler)
    return runResult
  })
}

addHook({
  name: 'tinypool',
  // version from tinypool@0.8 was used in vitest@1.6.0
  versions: ['>=0.8.0 <1.0.0'],
  file: 'dist/esm/index.js'
}, (TinyPool) => {
  wrapTinyPoolRun(TinyPool)
  return TinyPool
})

addHook({
  name: 'tinypool',
  versions: ['>=1.0.0'],
  file: 'dist/index.js'
}, (TinyPool) => {
  wrapTinyPoolRun(TinyPool)

  return TinyPool
})

function getWrappedOn (on) {
  return function (event, callback) {
    if (event !== 'message') {
      return on.apply(this, arguments)
    }
    // `arguments[1]` is the callback function, which
    // we modify to intercept our messages to not interfere
    // with vitest's own messages
    arguments[1] = shimmer.wrapFunction(callback, callback => function (message) {
      if (message.type !== 'Buffer' && Array.isArray(message)) {
        const [interprocessCode, data] = message
        if (interprocessCode === VITEST_WORKER_TRACE_PAYLOAD_CODE) {
          workerReportTraceCh.publish(data)
        } else if (interprocessCode === VITEST_WORKER_LOGS_PAYLOAD_CODE) {
          workerReportLogsCh.publish(data)
        }
        // If we execute the callback vitest crashes, as the message is not supported
        return
      }
      return callback.apply(this, arguments)
    })
    return on.apply(this, arguments)
  }
}

function getStartVitestWrapper (cliApiPackage, frameworkVersion) {
  if (!isCliApiPackage(cliApiPackage)) {
    return cliApiPackage
  }
  shimmer.wrap(cliApiPackage, 's', getCliOrStartVitestWrapper(frameworkVersion))

  if (hasForksPoolWorker(cliApiPackage)) {
    // function is async
    shimmer.wrap(cliApiPackage.f.prototype, 'start', start => function () {
      vitestPool = 'child_process'
      this.env.DD_VITEST_WORKER = '1'

      return start.apply(this, arguments)
    })
    shimmer.wrap(cliApiPackage.f.prototype, 'on', getWrappedOn)
  }

  if (hasThreadsPoolWorker(cliApiPackage)) {
    // function is async
    shimmer.wrap(cliApiPackage.T.prototype, 'start', start => function () {
      vitestPool = 'worker_threads'
      this.env.DD_VITEST_WORKER = '1'
      return start.apply(this, arguments)
    })
    shimmer.wrap(cliApiPackage.T.prototype, 'on', getWrappedOn)
  }
  return cliApiPackage
}

function wrapVitestTestRunner (VitestTestRunner) {
  // `onBeforeRunTask` is run before any repetition or attempt is run
  // `onBeforeRunTask` is an async function
  shimmer.wrap(VitestTestRunner.prototype, 'onBeforeRunTask', onBeforeRunTask => function (task) {
    const testName = getTestName(task)

    const {
      knownTests,
      isEarlyFlakeDetectionEnabled,
      isKnownTestsEnabled,
      numRepeats,
      isTestManagementTestsEnabled,
      testManagementAttemptToFixRetries,
      testManagementTests,
      isImpactedTestsEnabled,
      modifiedFiles
    } = getProvidedContext()

    if (isTestManagementTestsEnabled) {
      isAttemptToFixCh.publish({
        testManagementTests,
        testSuiteAbsolutePath: task.file.filepath,
        testName,
        onDone: (isAttemptToFix) => {
          if (isAttemptToFix) {
            isRetryReasonAttemptToFix = task.repeats !== testManagementAttemptToFixRetries
            task.repeats = testManagementAttemptToFixRetries
            attemptToFixTasks.add(task)
            taskToStatuses.set(task, [])
          }
        }
      })
      isDisabledCh.publish({
        testManagementTests,
        testSuiteAbsolutePath: task.file.filepath,
        testName,
        onDone: (isTestDisabled) => {
          if (isTestDisabled) {
            disabledTasks.add(task)
            if (!attemptToFixTasks.has(task)) {
              // we only actually skip if the test is not being attempted to be fixed
              task.mode = 'skip'
            }
          }
        }
      })
    }

    if (isImpactedTestsEnabled) {
      isModifiedCh.publish({
        modifiedFiles,
        testSuiteAbsolutePath: task.file.filepath,
        onDone: (isImpacted) => {
          if (isImpacted) {
            if (isEarlyFlakeDetectionEnabled) {
              isRetryReasonEfd = task.repeats !== numRepeats
              task.repeats = numRepeats
            }
            modifiedTasks.add(task)
            taskToStatuses.set(task, [])
          }
        }
      })
    }

    if (isKnownTestsEnabled) {
      isNewTestCh.publish({
        knownTests,
        testSuiteAbsolutePath: task.file.filepath,
        testName,
        onDone: (isNew) => {
          if (isNew && !attemptToFixTasks.has(task)) {
            if (isEarlyFlakeDetectionEnabled && !modifiedTasks.has(task)) {
              isRetryReasonEfd = task.repeats !== numRepeats
              task.repeats = numRepeats
            }
            newTasks.add(task)
            taskToStatuses.set(task, [])
          }
        }
      })
    }

    return onBeforeRunTask.apply(this, arguments)
  })

  // `onAfterRunTask` is run after all repetitions or attempts are run
  // `onAfterRunTask` is an async function
  shimmer.wrap(VitestTestRunner.prototype, 'onAfterRunTask', onAfterRunTask => function (task) {
    const { isEarlyFlakeDetectionEnabled, isTestManagementTestsEnabled } = getProvidedContext()

    if (isTestManagementTestsEnabled) {
      const isAttemptingToFix = attemptToFixTasks.has(task)
      const isDisabled = disabledTasks.has(task)
      const isQuarantined = quarantinedTasks.has(task)

      if (isAttemptingToFix && (isDisabled || isQuarantined)) {
        if (task.result.state === 'fail') {
          switchedStatuses.add(task)
        }
        task.result.state = 'pass'
      } else if (isQuarantined) {
        task.result.state = 'pass'
      }
    }

    if (isEarlyFlakeDetectionEnabled && taskToStatuses.has(task) && !attemptToFixTasks.has(task)) {
      const statuses = taskToStatuses.get(task)
      // If the test has passed at least once, we consider it passed
      if (statuses.includes('pass')) {
        if (task.result.state === 'fail') {
          switchedStatuses.add(task)
        }
        task.result.state = 'pass'
      }
    }

    return onAfterRunTask.apply(this, arguments)
  })

  // test start (only tests that are not marked as skip or todo)
  // `onBeforeTryTask` is run for every repetition and attempt of the test
  shimmer.wrap(VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task, retryInfo) {
    if (!testPassCh.hasSubscribers && !testErrorCh.hasSubscribers && !testSkipCh.hasSubscribers) {
      return onBeforeTryTask.apply(this, arguments)
    }
    const testName = getTestName(task)
    let isNew = false
    let isQuarantined = false

    const {
      isKnownTestsEnabled,
      isEarlyFlakeDetectionEnabled,
      isDiEnabled,
      isTestManagementTestsEnabled,
      testManagementTests,
      isFlakyTestRetriesEnabled
    } = getProvidedContext()

    if (isKnownTestsEnabled) {
      isNew = newTasks.has(task)
    }

    if (isTestManagementTestsEnabled) {
      isQuarantinedCh.publish({
        testManagementTests,
        testSuiteAbsolutePath: task.file.filepath,
        testName,
        onDone: (isTestQuarantined) => {
          isQuarantined = isTestQuarantined
          if (isTestQuarantined) {
            quarantinedTasks.add(task)
          }
        }
      })
    }

    const { retry: numAttempt, repeats: numRepetition } = retryInfo

    // We finish the previous test here because we know it has failed already
    if (numAttempt > 0) {
      const shouldWaitForHitProbe = isDiEnabled && numAttempt > 1
      if (shouldWaitForHitProbe) {
        await waitForHitProbe()
      }

      const promises = {}
      const shouldSetProbe = isDiEnabled && numAttempt === 1
      const ctx = taskToCtx.get(task)
      const testError = task.result?.errors?.[0]
      if (ctx) {
        testErrorCh.publish({
          error: testError,
          shouldSetProbe,
          promises,
          ...ctx.currentStore
        })
        // We wait for the probe to be set
        if (promises.setProbePromise) {
          await promises.setProbePromise
        }
      }
    }

    const lastExecutionStatus = task.result.state
    const shouldFlipStatus = isEarlyFlakeDetectionEnabled || attemptToFixTasks.has(task)
    const statuses = taskToStatuses.get(task)

    // These clauses handle task.repeats, whether EFD is enabled or not
    // The only thing that EFD does is to forcefully pass the test if it has passed at least once
    if (numRepetition > 0 && numRepetition < task.repeats) { // it may or may have not failed
      // Here we finish the earlier iteration,
      // as long as it's not the _last_ iteration (which will be finished normally)

      // TODO: check test duration (not to repeat if it's too slow)
      const ctx = taskToCtx.get(task)
      if (ctx) {
        if (lastExecutionStatus === 'fail') {
          const testError = task.result?.errors?.[0]
          testErrorCh.publish({ error: testError, ...ctx.currentStore })
        } else {
          testPassCh.publish({ task, ...ctx.currentStore })
        }
        if (shouldFlipStatus) {
          statuses.push(lastExecutionStatus)
          // 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:
          // if the last attempt passes, vitest would consider the test as failed
          // but after this change, it will consider the test as passed
          task.result.state = 'pass'
        }
      }
    } else if (numRepetition === task.repeats) {
      if (shouldFlipStatus) {
        statuses.push(lastExecutionStatus)
      }

      const ctx = taskToCtx.get(task)
      if (lastExecutionStatus === 'fail') {
        const testError = task.result?.errors?.[0]
        testErrorCh.publish({ error: testError, ...ctx.currentStore })
      } else {
        testPassCh.publish({ task, ...ctx.currentStore })
      }
    }

    const isRetryReasonAtr = numAttempt > 0 &&
      isFlakyTestRetriesEnabled &&
      !isRetryReasonAttemptToFix &&
      !isRetryReasonEfd

    const ctx = {
      testName,
      testSuiteAbsolutePath: task.file.filepath,
      isRetry: numAttempt > 0 || numRepetition > 0,
      isRetryReasonEfd,
      isRetryReasonAttemptToFix: isRetryReasonAttemptToFix && numRepetition > 0,
      isNew,
      mightHitProbe: isDiEnabled && numAttempt > 0,
      isAttemptToFix: attemptToFixTasks.has(task),
      isDisabled: disabledTasks.has(task),
      isQuarantined,
      isRetryReasonAtr,
      isModified: modifiedTasks.has(task)
    }
    taskToCtx.set(task, ctx)

    testStartCh.runStores(ctx, () => {})
    return onBeforeTryTask.apply(this, arguments)
  })

  // test finish (only passed tests)
  shimmer.wrap(VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask =>
    async function (task, { retry: retryCount }) {
      if (!testPassCh.hasSubscribers && !testErrorCh.hasSubscribers && !testSkipCh.hasSubscribers) {
        return onAfterTryTask.apply(this, arguments)
      }
      const result = await onAfterTryTask.apply(this, arguments)

      const { testManagementAttemptToFixRetries } = getProvidedContext()

      const status = getVitestTestStatus(task, retryCount)
      const ctx = taskToCtx.get(task)

      const { isDiEnabled } = getProvidedContext()

      if (isDiEnabled && retryCount > 1) {
        await waitForHitProbe()
      }

      let attemptToFixPassed = false
      let attemptToFixFailed = false
      if (attemptToFixTasks.has(task)) {
        const statuses = taskToStatuses.get(task)
        if (statuses.length === testManagementAttemptToFixRetries) {
          if (statuses.every(status => status === 'pass')) {
            attemptToFixPassed = true
          } else if (statuses.includes('fail')) {
            attemptToFixFailed = true
          }
        }
      }

      if (ctx) {
        // We don't finish here because the test might fail in a later hook (afterEach)
        ctx.status = status
        ctx.task = task
        ctx.attemptToFixPassed = attemptToFixPassed
        ctx.attemptToFixFailed = attemptToFixFailed
        testFinishTimeCh.runStores(ctx, () => {})
      }

      return result
    })
}

addHook({
  name: 'vitest',
  versions: ['>=4.0.0'],
  filePattern: 'dist/chunks/test.*'
}, (testPackage) => {
  if (!isTestPackage(testPackage)) {
    return testPackage
  }

  wrapVitestTestRunner(testPackage.V)

  return testPackage
})

addHook({
  name: 'vitest',
  versions: ['>=1.6.0 <4.0.0'],
  file: 'dist/runners.js'
}, (vitestPackage) => {
  const { VitestTestRunner } = vitestPackage

  wrapVitestTestRunner(VitestTestRunner)

  return vitestPackage
})

// There are multiple index* files across different versions of vitest,
// so we check for the existence of BaseSequencer to determine if we are in the right file
addHook({
  name: 'vitest',
  versions: ['>=1.6.0 <2.0.0'],
  filePattern: 'dist/vendor/index.*'
}, (vitestPackage) => {
  if (isReporterPackage(vitestPackage)) {
    shimmer.wrap(vitestPackage.B.prototype, 'sort', getSortWrapper)
  }

  return vitestPackage
})

addHook({
  name: 'vitest',
  versions: ['>=2.0.0 <2.0.5'],
  filePattern: 'dist/vendor/index.*'
}, (vitestPackage) => {
  if (isReporterPackageNew(vitestPackage)) {
    shimmer.wrap(vitestPackage.e.prototype, 'sort', getSortWrapper)
  }

  return vitestPackage
})

addHook({
  name: 'vitest',
  versions: ['>=2.0.5 <2.1.0'],
  filePattern: 'dist/chunks/index.*'
}, (vitestPackage) => {
  if (isReporterPackageNewest(vitestPackage)) {
    shimmer.wrap(vitestPackage.h.prototype, 'sort', getSortWrapper)
  }

  return vitestPackage
})

addHook({
  name: 'vitest',
  versions: ['>=2.1.0 <3.0.0'],
  filePattern: 'dist/chunks/RandomSequencer.*'
}, (randomSequencerPackage) => {
  shimmer.wrap(randomSequencerPackage.B.prototype, 'sort', getSortWrapper)
  return randomSequencerPackage
})

addHook({
  name: 'vitest',
  versions: ['>=3.0.9'],
  filePattern: 'dist/chunks/coverage.*'
}, (coveragePackage) => {
  if (isBaseSequencer(coveragePackage)) {
    shimmer.wrap(coveragePackage.b.prototype, 'sort', getSortWrapper)
  }
  return coveragePackage
})

addHook({
  name: 'vitest',
  versions: ['>=3.0.0 <3.0.9'],
  filePattern: 'dist/chunks/resolveConfig.*'
}, (resolveConfigPackage) => {
  shimmer.wrap(resolveConfigPackage.B.prototype, 'sort', getSortWrapper)
  return resolveConfigPackage
})

// Can't specify file because compiled vitest includes hashes in their files
// Following 3 wrappers are for test session start
addHook({
  name: 'vitest',
  versions: ['>=1.6.0 <2.0.5'],
  filePattern: 'dist/vendor/cac.*'
}, getCreateCliWrapper)

addHook({
  name: 'vitest',
  versions: ['>=2.0.5'],
  filePattern: 'dist/chunks/cac.*'
}, getCreateCliWrapper)

addHook({
  name: 'vitest',
  versions: ['>=1.6.0 <2.0.5'],
  filePattern: 'dist/vendor/cli-api.*'
}, getStartVitestWrapper)

addHook({
  name: 'vitest',
  versions: ['>=2.0.5'],
  filePattern: 'dist/chunks/cli-api.*'
}, getStartVitestWrapper)

// test suite start and finish
// only relevant for workers
addHook({
  name: '@vitest/runner',
  versions: ['>=1.6.0'],
  file: 'dist/index.js'
}, (vitestPackage, frameworkVersion) => {
  shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPaths) {
    let testSuiteError = null
    if (!testSuiteFinishCh.hasSubscribers) {
      return startTests.apply(this, arguments)
    }
    // From >=3.0.1, the first arguments changes from a string to an object containing the filepath
    const testSuiteAbsolutePath = testPaths[0]?.filepath || testPaths[0]

    const testSuiteCtx = { testSuiteAbsolutePath, frameworkVersion }
    testSuiteStartCh.runStores(testSuiteCtx, () => {})
    const startTestsResponse = await startTests.apply(this, arguments)

    let onFinish = null
    const onFinishPromise = new Promise(resolve => {
      onFinish = resolve
    })

    const testTasks = getTypeTasks(startTestsResponse[0].tasks)

    // Only one test task per test, even if there are retries
    testTasks.forEach(task => {
      const testCtx = taskToCtx.get(task)
      const { result } = task
      // We have to trick vitest into thinking that the test has passed
      // but we want to report it as failed if it did fail
      const isSwitchedStatus = switchedStatuses.has(task)

      if (result) {
        const { state, duration, errors } = result
        if (state === 'skip') { // programmatic skip
          testSkipCh.publish({
            testName: getTestName(task),
            testSuiteAbsolutePath: task.file.filepath,
            isNew: newTasks.has(task),
            isDisabled: disabledTasks.has(task)
          })
        } else if (state === 'pass' && !isSwitchedStatus) {
          if (testCtx) {
            testPassCh.publish({ task, ...testCtx.currentStore })
          }
        } else if (state === 'fail' || isSwitchedStatus) {
          let testError

          if (errors?.length) {
            testError = errors[0]
          }

          let hasFailedAllRetries = false
          let attemptToFixFailed = false
          if (attemptToFixTasks.has(task)) {
            const statuses = taskToStatuses.get(task)
            if (statuses.includes('fail')) {
              attemptToFixFailed = true
            }
            if (statuses.every(status => status === 'fail')) {
              hasFailedAllRetries = true
            }
          }

          if (testCtx) {
            const isRetry = task.result?.retryCount > 0
            // `duration` is the duration of all the retries, so it can't be used if there are retries
            testErrorCh.publish({
              duration: isRetry ? undefined : duration,
              error: testError,
              hasFailedAllRetries,
              attemptToFixFailed,
              ...testCtx.currentStore
            })
          }
          if (errors?.length) {
            testSuiteError = testError // we store the error to bubble it up to the suite
          }
        }
      } else { // test.skip or test.todo
        testSkipCh.publish({
          testName: getTestName(task),
          testSuiteAbsolutePath: task.file.filepath,
          isNew: newTasks.has(task),
          isDisabled: disabledTasks.has(task)
        })
      }
    })

    const testSuiteResult = startTestsResponse[0].result

    if (testSuiteResult.errors?.length) { // Errors from root level hooks
      testSuiteError = testSuiteResult.errors[0]
    } else if (testSuiteResult.state === 'fail') { // Errors from `describe` level hooks
      const suiteTasks = getTypeTasks(startTestsResponse[0].tasks, 'suite')
      const failedSuites = suiteTasks.filter(task => task.result?.state === 'fail')
      if (failedSuites.length && failedSuites[0].result?.errors?.length) {
        testSuiteError = failedSuites[0].result.errors[0]
      }
    }

    if (testSuiteError) {
      testSuiteCtx.error = testSuiteError
      testSuiteErrorCh.runStores(testSuiteCtx, () => {})
    }

    testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish, ...testSuiteCtx.currentStore })

    await onFinishPromise

    return startTestsResponse
  })

  return vitestPackage
})