Skip to content

Commit

Permalink
fix vitest, now for real
Browse files Browse the repository at this point in the history
  • Loading branch information
juan-fernandez committed Jan 10, 2025
1 parent 6cce993 commit 410c980
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 68 deletions.
7 changes: 2 additions & 5 deletions integration-tests/vitest/vitest.spec.js
Original file line number Diff line number Diff line change
Expand Up @@ -904,7 +904,7 @@ versions.forEach((version) => {

// dynamic instrumentation only supported from >=2.0.0
if (version === 'latest') {
context.only('dynamic instrumentation', () => {
context('dynamic instrumentation', () => {
it('does not activate it if DD_TEST_DYNAMIC_INSTRUMENTATION_ENABLED is not set', (done) => {
receiver.setSettings({
flaky_test_retries_enabled: true,
Expand Down Expand Up @@ -1007,7 +1007,7 @@ versions.forEach((version) => {
})
})

it.only('runs retries with dynamic instrumentation', (done) => {
it('runs retries with dynamic instrumentation', (done) => {
receiver.setSettings({
flaky_test_retries_enabled: true,
di_enabled: true
Expand Down Expand Up @@ -1087,9 +1087,6 @@ versions.forEach((version) => {
}
)

childProcess.stdout.pipe(process.stdout)
childProcess.stderr.pipe(process.stderr)

childProcess.on('exit', () => {
Promise.all([eventsPromise, logsPromise]).then(() => {
assert.equal(snapshotIdByTest, snapshotIdByLog)
Expand Down
129 changes: 82 additions & 47 deletions packages/datadog-instrumentations/src/vitest.js
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,42 @@ const newTasks = new WeakSet()
const switchedStatuses = new WeakSet()
const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')

const BREAKPOINT_HIT_GRACE_PERIOD_MS = 400

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

function getProvidedContext () {
try {
const {
_ddIsEarlyFlakeDetectionEnabled,
_ddIsDiEnabled,
_ddKnownTests: knownTests,
_ddEarlyFlakeDetectionNumRetries: numRepeats
} = globalThis.__vitest_worker__.providedContext

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

function isReporterPackage (vitestPackage) {
return vitestPackage.B?.name === 'BaseSequencer'
}
Expand Down Expand Up @@ -253,39 +289,34 @@ addHook({
// `onBeforeRunTask` is run before any repetition or attempt is run
shimmer.wrap(VitestTestRunner.prototype, 'onBeforeRunTask', onBeforeRunTask => async function (task) {
const testName = getTestName(task)
try {
const {
_ddKnownTests: knownTests,
_ddIsEarlyFlakeDetectionEnabled: isEarlyFlakeDetectionEnabled,
_ddEarlyFlakeDetectionNumRetries: numRepeats
} = globalThis.__vitest_worker__.providedContext

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

const {
_ddKnownTests: knownTests,
_ddIsEarlyFlakeDetectionEnabled: isEarlyFlakeDetectionEnabled,
_ddEarlyFlakeDetectionNumRetries: numRepeats
} = getProvidedContext()

if (isEarlyFlakeDetectionEnabled) {
isNewTestCh.publish({
knownTests,
testSuiteAbsolutePath: task.file.filepath,
testName,
onDone: (isNew) => {
if (isNew) {
task.repeats = numRepeats
newTasks.add(task)
taskToStatuses.set(task, [])
}
})
}
} catch (e) {
log.error('Vitest workers could not parse known tests, so Early Flake Detection will not work.')
}
})
}

return onBeforeRunTask.apply(this, arguments)
})

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

if (isEarlyFlakeDetectionEnabled && taskToStatuses.has(task)) {
const statuses = taskToStatuses.get(task)
Expand All @@ -309,43 +340,40 @@ addHook({
}
const testName = getTestName(task)
let isNew = false
let isEarlyFlakeDetectionEnabled = false
let isDiEnabled = false

try {
const {
_ddIsEarlyFlakeDetectionEnabled,
_ddIsDiEnabled
} = globalThis.__vitest_worker__.providedContext

isEarlyFlakeDetectionEnabled = _ddIsEarlyFlakeDetectionEnabled
isDiEnabled = _ddIsDiEnabled
const {
isEarlyFlakeDetectionEnabled,
isDiEnabled
} = getProvidedContext()

if (isEarlyFlakeDetectionEnabled) {
isNew = newTasks.has(task)
}
} catch (e) {
log.error('Vitest workers could not parse known tests, so Early Flake Detection will not work.')
if (isEarlyFlakeDetectionEnabled) {
isNew = newTasks.has(task)
}

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

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

const promises = {}
const shouldSetProbe = isDiEnabled && numAttempt === 1
const asyncResource = taskToAsync.get(task)
const testError = task.result?.errors?.[0]
if (asyncResource) {
asyncResource.runInAsyncScope(() => {
testErrorCh.publish({
error: testError,
willBeRetried: true,
probe,
isDiEnabled
shouldSetProbe,
promises
})
})
// We wait for the probe to be set
if (probe.setProbePromise) {
await probe.setProbePromise
if (promises.setProbePromise) {
await promises.setProbePromise
}
}
}
Expand Down Expand Up @@ -401,7 +429,8 @@ addHook({
testName,
testSuiteAbsolutePath: task.file.filepath,
isRetry: numAttempt > 0 || numRepetition > 0,
isNew
isNew,
mightHitProbe: isDiEnabled && numAttempt > 0
})
})
return onBeforeTryTask.apply(this, arguments)
Expand All @@ -418,6 +447,12 @@ addHook({
const status = getVitestTestStatus(task, retryCount)
const asyncResource = taskToAsync.get(task)

const { isDiEnabled } = getProvidedContext()

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

if (asyncResource) {
// We don't finish here because the test might fail in a later hook (afterEach)
asyncResource.runInAsyncScope(() => {
Expand Down
28 changes: 12 additions & 16 deletions packages/datadog-plugin-vitest/src/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,7 @@ const {
TEST_SOURCE_START,
TEST_IS_NEW,
TEST_EARLY_FLAKE_ENABLED,
TEST_EARLY_FLAKE_ABORT_REASON,
TEST_NAME,
DI_ERROR_DEBUG_INFO_CAPTURED,
DI_DEBUG_ERROR_SNAPSHOT_ID,
DI_DEBUG_ERROR_FILE,
DI_DEBUG_ERROR_LINE
TEST_EARLY_FLAKE_ABORT_REASON
} = require('../../dd-trace/src/plugins/util/test')
const { COMPONENT } = require('../../dd-trace/src/constants')
const {
Expand Down Expand Up @@ -65,8 +60,7 @@ class VitestPlugin extends CiPlugin {
onDone(isFaulty)
})

this.addSub('ci:vitest:test:start', ({ testName, testSuiteAbsolutePath, isRetry, isNew }) => {
console.log('test:start', { isRetry, testName })
this.addSub('ci:vitest:test:start', ({ testName, testSuiteAbsolutePath, isRetry, isNew, mightHitProbe }) => {
const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
const store = storage.getStore()

Expand All @@ -88,7 +82,12 @@ class VitestPlugin extends CiPlugin {
)

this.enter(span, store)
this.activeTestSpan = span

// TODO: there might be multiple tests for which mightHitProbe is true, so activeTestSpan
// might be wrongly overwritten.
if (mightHitProbe) {
this.activeTestSpan = span
}
})

this.addSub('ci:vitest:test:finish-time', ({ status, task }) => {
Expand All @@ -114,23 +113,21 @@ class VitestPlugin extends CiPlugin {
span.setTag(TEST_STATUS, 'pass')
span.finish(this.taskToFinishTime.get(task))
finishAllTraceSpans(span)
this.activeTestSpan = null
}
})

this.addSub('ci:vitest:test:error', ({ duration, error, willBeRetried, probe, isDiEnabled }) => {
this.addSub('ci:vitest:test:error', ({ duration, error, shouldSetProbe, promises }) => {
const store = storage.getStore()
const span = store?.span

if (span) {
// TODO: PROBE SHOULD ONLY BE ADDED IN THE FIRST TRY!
if (willBeRetried && this.di && isDiEnabled) {
if (shouldSetProbe && this.di) {
const probeInformation = this.addDiProbe(error)
if (probeInformation) {
const { probeId, stackIndex, setProbePromise } = probeInformation
this.runningTestProbeId = probeId
this.testErrorStackIndex = stackIndex
probe.setProbePromise = setProbePromise
promises.setProbePromise = setProbePromise
}
}
this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', {
Expand All @@ -144,10 +141,9 @@ class VitestPlugin extends CiPlugin {
if (duration) {
span.finish(span._startTime + duration - MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION) // milliseconds
} else {
span.finish() // retries will not have a duration
span.finish() // `duration` is empty for retries, so we'll use clock time
}
finishAllTraceSpans(span)
this.activeTestSpan = null
}
})

Expand Down

0 comments on commit 410c980

Please sign in to comment.