From 821b597a184ba1e1a0a03dcffda406725c4a7ca3 Mon Sep 17 00:00:00 2001
From: AriPerkkio <ari.perkkio@gmail.com>
Date: Tue, 18 Apr 2023 15:20:41 +0300
Subject: [PATCH] feat: `--bail` option for cancelling test run

---
 docs/config/index.md                       |  8 ++++
 package.json                               |  2 +-
 packages/browser/src/client/runner.ts      | 12 +++++-
 packages/runner/src/types/runner.ts        |  2 +-
 packages/vitest/src/api/setup.ts           |  6 +++
 packages/vitest/src/api/types.ts           |  2 +
 packages/vitest/src/node/cli.ts            |  1 +
 packages/vitest/src/node/pools/rpc.ts      |  6 +++
 packages/vitest/src/node/reporters/base.ts |  1 +
 packages/vitest/src/node/state.ts          |  8 +++-
 packages/vitest/src/runtime/child.ts       |  2 +-
 packages/vitest/src/runtime/entry.ts       | 14 +++++++
 packages/vitest/src/runtime/worker.ts      |  2 +-
 packages/vitest/src/types/config.ts        |  5 +++
 packages/vitest/src/types/rpc.ts           |  2 +
 pnpm-lock.yaml                             | 18 +++++++++
 test/bail/fixtures/test/first.test.ts      | 14 +++++++
 test/bail/fixtures/test/second.test.ts     | 17 +++++++++
 test/bail/fixtures/vitest.config.ts        | 35 +++++++++++++++++
 test/bail/package.json                     | 14 +++++++
 test/bail/test/bail.test.ts                | 44 ++++++++++++++++++++++
 test/bail/vitest.config.ts                 | 14 +++++++
 22 files changed, 223 insertions(+), 6 deletions(-)
 create mode 100644 test/bail/fixtures/test/first.test.ts
 create mode 100644 test/bail/fixtures/test/second.test.ts
 create mode 100644 test/bail/fixtures/vitest.config.ts
 create mode 100644 test/bail/package.json
 create mode 100644 test/bail/test/bail.test.ts
 create mode 100644 test/bail/vitest.config.ts

diff --git a/docs/config/index.md b/docs/config/index.md
index 8951feafdc7b..43be6387a4f8 100644
--- a/docs/config/index.md
+++ b/docs/config/index.md
@@ -1375,3 +1375,11 @@ Influences whether or not the `showDiff` flag should be included in the thrown A
 Sets length threshold for actual and expected values in assertion errors. If this threshold is exceeded, for example for large data structures, the value is replaced with something like `[ Array(3) ]` or `{ Object (prop1, prop2) }`. Set it to `0` if you want to disable truncating altogether.
 
 This config option affects truncating values in `test.each` titles and inside the assertion error message.
+
+### bail
+
+- **Type:** `number`
+- **Default:** `0`
+- **CLI**: `--bail=<value>`
+
+Stop test execution when given number of tests have failed.
diff --git a/package.json b/package.json
index 713ed041f43a..2c23bf7c34d0 100644
--- a/package.json
+++ b/package.json
@@ -23,7 +23,7 @@
     "test:run": "vitest run -r test/core",
     "test:all": "CI=true pnpm -r --stream run test --allowOnly",
     "test:ci": "CI=true pnpm -r --stream --filter !test-fails --filter !test-browser --filter !test-esm --filter !test-browser run test --allowOnly",
-    "test:ci:single-thread": "CI=true pnpm -r --stream --filter !test-fails --filter !test-coverage --filter !test-watch --filter !test-esm --filter !test-browser run test --allowOnly --no-threads",
+    "test:ci:single-thread": "CI=true pnpm -r --stream --filter !test-fails --filter !test-coverage --filter !test-watch --filter !test-bail --filter !test-esm --filter !test-browser run test --allowOnly --no-threads",
     "typecheck": "tsc --noEmit",
     "typecheck:why": "tsc --noEmit --explainFiles > explainTypes.txt",
     "ui:build": "vite build packages/ui",
diff --git a/packages/browser/src/client/runner.ts b/packages/browser/src/client/runner.ts
index 3bdc40c1f0d2..d73c2ccc6ea3 100644
--- a/packages/browser/src/client/runner.ts
+++ b/packages/browser/src/client/runner.ts
@@ -23,10 +23,20 @@ export function createBrowserRunner(original: any, coverageModule: CoverageHandl
     }
 
     async onAfterRunTest(task: Test) {
-      await super.onAfterRunTest?.()
+      await super.onAfterRunTest?.(task)
       task.result?.errors?.forEach((error) => {
         console.error(error.message)
       })
+
+      if (this.config.bail && task.result?.state === 'fail') {
+        const previousFailures = await rpc().getCountOfFailedTests()
+        const currentFailures = 1 + previousFailures
+
+        if (currentFailures >= this.config.bail) {
+          rpc().onCancel('test-failure')
+          this.onCancel?.('test-failure')
+        }
+      }
     }
 
     async onAfterRunSuite() {
diff --git a/packages/runner/src/types/runner.ts b/packages/runner/src/types/runner.ts
index e74b93dcbb36..694d6c70045d 100644
--- a/packages/runner/src/types/runner.ts
+++ b/packages/runner/src/types/runner.ts
@@ -27,7 +27,7 @@ export interface VitestRunnerConstructor {
   new(config: VitestRunnerConfig): VitestRunner
 }
 
-export type CancelReason = 'keyboard-input' | string & {}
+export type CancelReason = 'keyboard-input' | 'test-failure' | string & {}
 
 export interface VitestRunner {
   /**
diff --git a/packages/vitest/src/api/setup.ts b/packages/vitest/src/api/setup.ts
index 6fe9999ca4cb..545ff2cf00bc 100644
--- a/packages/vitest/src/api/setup.ts
+++ b/packages/vitest/src/api/setup.ts
@@ -111,6 +111,12 @@ export function setup(vitestOrWorkspace: Vitest | WorkspaceProject, server?: Vit
             return ctx.updateSnapshot()
           return ctx.updateSnapshot([file.filepath])
         },
+        onCancel(reason) {
+          ctx.cancelCurrentRun(reason)
+        },
+        getCountOfFailedTests() {
+          return ctx.state.getCountOfFailedTests()
+        },
       },
       {
         post: msg => ws.send(msg),
diff --git a/packages/vitest/src/api/types.ts b/packages/vitest/src/api/types.ts
index e3b0dcac024f..3f96018b2d82 100644
--- a/packages/vitest/src/api/types.ts
+++ b/packages/vitest/src/api/types.ts
@@ -11,6 +11,8 @@ export interface WebSocketHandlers {
   onTaskUpdate(packs: TaskResultPack[]): void
   onAfterSuiteRun(meta: AfterSuiteRunMeta): void
   onDone(name: string): void
+  onCancel(reason: CancelReason): void
+  getCountOfFailedTests(): number
   sendLog(log: UserConsoleLog): void
   getFiles(): File[]
   getPaths(): string[]
diff --git a/packages/vitest/src/node/cli.ts b/packages/vitest/src/node/cli.ts
index 617b20494379..f9599ad667fe 100644
--- a/packages/vitest/src/node/cli.ts
+++ b/packages/vitest/src/node/cli.ts
@@ -45,6 +45,7 @@ cli
   .option('--inspect', 'Enable Node.js inspector')
   .option('--inspect-brk', 'Enable Node.js inspector with break')
   .option('--test-timeout <time>', 'Default timeout of a test in milliseconds (default: 5000)')
+  .option('--bail <number>', 'Stop test execution when given number of tests have failed', { default: 0 })
   .help()
 
 cli
diff --git a/packages/vitest/src/node/pools/rpc.ts b/packages/vitest/src/node/pools/rpc.ts
index 1e295f9a6576..0b67d50e4b8a 100644
--- a/packages/vitest/src/node/pools/rpc.ts
+++ b/packages/vitest/src/node/pools/rpc.ts
@@ -58,5 +58,11 @@ export function createMethodsRPC(project: WorkspaceProject): RuntimeRPC {
     onFinished(files) {
       project.report('onFinished', files, ctx.state.getUnhandledErrors())
     },
+    onCancel(reason) {
+      ctx.cancelCurrentRun(reason)
+    },
+    getCountOfFailedTests() {
+      return ctx.state.getCountOfFailedTests()
+    },
   }
 }
diff --git a/packages/vitest/src/node/reporters/base.ts b/packages/vitest/src/node/reporters/base.ts
index 484840208bec..2a55523f19ee 100644
--- a/packages/vitest/src/node/reporters/base.ts
+++ b/packages/vitest/src/node/reporters/base.ts
@@ -50,6 +50,7 @@ export abstract class BaseReporter implements Reporter {
 
   async onFinished(files = this.ctx.state.getFiles(), errors = this.ctx.state.getUnhandledErrors()) {
     this.end = performance.now()
+
     await this.reportSummary(files)
     if (errors.length) {
       if (!this.ctx.config.dangerouslyIgnoreUnhandledErrors)
diff --git a/packages/vitest/src/node/state.ts b/packages/vitest/src/node/state.ts
index bbed9381447d..c2ef24e81f18 100644
--- a/packages/vitest/src/node/state.ts
+++ b/packages/vitest/src/node/state.ts
@@ -127,6 +127,10 @@ export class StateManager {
     }
   }
 
+  getCountOfFailedTests() {
+    return Array.from(this.idMap.values()).filter(t => t.result?.state === 'fail').length
+  }
+
   cancelFiles(files: string[], root: string) {
     this.collectFiles(files.map(filepath => ({
       filepath,
@@ -134,7 +138,9 @@ export class StateManager {
       id: filepath,
       mode: 'skip',
       type: 'suite',
-
+      result: {
+        state: 'skip',
+      },
       // Cancelled files have not yet collected tests
       tasks: [],
     })))
diff --git a/packages/vitest/src/runtime/child.ts b/packages/vitest/src/runtime/child.ts
index 9240b20226e7..7e606f7f4a45 100644
--- a/packages/vitest/src/runtime/child.ts
+++ b/packages/vitest/src/runtime/child.ts
@@ -38,7 +38,7 @@ function init(ctx: ChildContext) {
         onCancel: setCancel,
       },
       {
-        eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'],
+        eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit', 'onCancel'],
         serialize: v8.serialize,
         deserialize: v => v8.deserialize(Buffer.from(v)),
         post(v) {
diff --git a/packages/vitest/src/runtime/entry.ts b/packages/vitest/src/runtime/entry.ts
index 5c42e497ab16..2c16e82e58b1 100644
--- a/packages/vitest/src/runtime/entry.ts
+++ b/packages/vitest/src/runtime/entry.ts
@@ -71,6 +71,20 @@ async function getTestRunner(config: ResolvedConfig, executor: VitestExecutor):
     await originalOnAfterRun?.call(testRunner, files)
   }
 
+  const originalOnAfterRunTest = testRunner.onAfterRunTest
+  testRunner.onAfterRunTest = async (test) => {
+    if (config.bail && test.result?.state === 'fail') {
+      const previousFailures = await rpc().getCountOfFailedTests()
+      const currentFailures = 1 + previousFailures
+
+      if (currentFailures >= config.bail) {
+        rpc().onCancel('test-failure')
+        testRunner.onCancel?.('test-failure')
+      }
+    }
+    await originalOnAfterRunTest?.call(testRunner, test)
+  }
+
   return testRunner
 }
 
diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts
index 983303333931..0db3280c364a 100644
--- a/packages/vitest/src/runtime/worker.ts
+++ b/packages/vitest/src/runtime/worker.ts
@@ -41,7 +41,7 @@ function init(ctx: WorkerContext) {
         onCancel: setCancel,
       },
       {
-        eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'],
+        eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit', 'onCancel'],
         post(v) { port.postMessage(v) },
         on(fn) { port.addListener('message', fn) },
       },
diff --git a/packages/vitest/src/types/config.ts b/packages/vitest/src/types/config.ts
index 4b6464c0c933..858303728c59 100644
--- a/packages/vitest/src/types/config.ts
+++ b/packages/vitest/src/types/config.ts
@@ -567,6 +567,11 @@ export interface InlineConfig {
    * https://github.com/chaijs/chai/blob/4.x.x/lib/chai/config.js
   */
   chaiConfig?: ChaiConfig
+
+  /**
+   * Stop test execution when given number of tests have failed.
+   */
+  bail?: number
 }
 
 export interface TypecheckConfig {
diff --git a/packages/vitest/src/types/rpc.ts b/packages/vitest/src/types/rpc.ts
index cd6b0402e7e8..bcfd6322fe11 100644
--- a/packages/vitest/src/types/rpc.ts
+++ b/packages/vitest/src/types/rpc.ts
@@ -20,6 +20,8 @@ export interface RuntimeRPC {
   onCollected: (files: File[]) => void
   onAfterSuiteRun: (meta: AfterSuiteRunMeta) => void
   onTaskUpdate: (pack: TaskResultPack[]) => void
+  onCancel(reason: CancelReason): void
+  getCountOfFailedTests(): number
 
   snapshotSaved: (snapshot: SnapshotResult) => void
   resolveSnapshotPath: (testPath: string) => string
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index e9d8e9c5699c..c505cddc9f0c 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1375,6 +1375,24 @@ importers:
         specifier: ^2.79.1
         version: 2.79.1
 
+  test/bail:
+    devDependencies:
+      '@vitest/browser':
+        specifier: workspace:*
+        version: link:../../packages/browser
+      execa:
+        specifier: ^7.0.0
+        version: 7.1.1
+      vite:
+        specifier: ^4.2.1
+        version: 4.2.1(@types/node@18.15.11)
+      vitest:
+        specifier: workspace:*
+        version: link:../../packages/vitest
+      webdriverio:
+        specifier: latest
+        version: 8.8.0(typescript@5.0.3)
+
   test/base:
     devDependencies:
       vitest:
diff --git a/test/bail/fixtures/test/first.test.ts b/test/bail/fixtures/test/first.test.ts
new file mode 100644
index 000000000000..3e4940e240f3
--- /dev/null
+++ b/test/bail/fixtures/test/first.test.ts
@@ -0,0 +1,14 @@
+import { expect, test } from 'vitest'
+
+test('1 - first.test.ts - this should pass', async () => {
+  await new Promise(resolve => setTimeout(resolve, 250))
+  expect(true).toBeTruthy()
+})
+
+test('2 - first.test.ts - this should fail', () => {
+  expect(false).toBeTruthy()
+})
+
+test('3 - first.test.ts - this should be skipped', () => {
+  expect(true).toBeTruthy()
+})
diff --git a/test/bail/fixtures/test/second.test.ts b/test/bail/fixtures/test/second.test.ts
new file mode 100644
index 000000000000..bcb2d28aab09
--- /dev/null
+++ b/test/bail/fixtures/test/second.test.ts
@@ -0,0 +1,17 @@
+import { expect, test } from 'vitest'
+
+// When using multi threads the first test will start before failing.test.ts fails
+const isThreads = process.env.THREADS === 'true'
+
+test(`1 - second.test.ts - this should ${isThreads ? 'pass' : 'be skipped'}`, async () => {
+  await new Promise(resolve => setTimeout(resolve, 1500))
+  expect(true).toBeTruthy()
+})
+
+test('2 - second.test.ts - this should be skipped', () => {
+  expect(true).toBeTruthy()
+})
+
+test('3 - second.test.ts - this should be skipped', () => {
+  expect(true).toBeTruthy()
+})
diff --git a/test/bail/fixtures/vitest.config.ts b/test/bail/fixtures/vitest.config.ts
new file mode 100644
index 000000000000..7482dfe640d6
--- /dev/null
+++ b/test/bail/fixtures/vitest.config.ts
@@ -0,0 +1,35 @@
+import { defineConfig } from 'vitest/config'
+import type { WorkspaceSpec } from 'vitest/node'
+
+class TestNameSequencer {
+  async sort(files: WorkspaceSpec[]): Promise<WorkspaceSpec[]> {
+    return [...files].sort(([, filenameA], [, filenameB]) => {
+      if (filenameA > filenameB)
+        return 1
+
+      if (filenameA < filenameB)
+        return -1
+
+      return 0
+    })
+  }
+
+  public async shard(files: WorkspaceSpec[]): Promise<WorkspaceSpec[]> {
+    return files
+  }
+}
+
+export default defineConfig({
+  test: {
+    reporters: 'verbose',
+    cache: false,
+    watch: false,
+    sequence: {
+      sequencer: TestNameSequencer,
+    },
+    browser: {
+      headless: true,
+      name: 'chrome',
+    },
+  },
+})
diff --git a/test/bail/package.json b/test/bail/package.json
new file mode 100644
index 000000000000..694a5df093dc
--- /dev/null
+++ b/test/bail/package.json
@@ -0,0 +1,14 @@
+{
+  "name": "@vitest/test-bail",
+  "private": true,
+  "scripts": {
+    "test": "vitest"
+  },
+  "devDependencies": {
+    "@vitest/browser": "workspace:*",
+    "execa": "^7.0.0",
+    "vite": "latest",
+    "vitest": "workspace:*",
+    "webdriverio": "latest"
+  }
+}
diff --git a/test/bail/test/bail.test.ts b/test/bail/test/bail.test.ts
new file mode 100644
index 000000000000..f42293710a14
--- /dev/null
+++ b/test/bail/test/bail.test.ts
@@ -0,0 +1,44 @@
+import { expect, test } from 'vitest'
+import { execa } from 'execa'
+
+const configs: string[][] = []
+const pools = [['--threads', 'true'], ['--threads', 'false'], ['--single-thread']]
+
+if (process.platform !== 'win32')
+  pools.push(['--browser'])
+
+for (const isolate of ['true', 'false']) {
+  for (const pool of pools) {
+    configs.push([
+      '--bail',
+      '1',
+      '--isolate',
+      isolate,
+      ...pool,
+    ])
+  }
+}
+
+for (const config of configs) {
+  test(`should bail with "${config.join(' ')}"`, async () => {
+    const { exitCode, stdout } = await execa('vitest', [
+      '--no-color',
+      '--root',
+      'fixtures',
+      ...config,
+    ], {
+      reject: false,
+      env: { THREADS: config.join(' ').includes('--threads true') ? 'true' : 'false' },
+    })
+
+    expect(exitCode).toBe(1)
+    expect(stdout).toMatch('✓ test/first.test.ts > 1 - first.test.ts - this should pass')
+    expect(stdout).toMatch('× test/first.test.ts > 2 - first.test.ts - this should fail')
+
+    // Cancelled tests should not be run
+    expect(stdout).not.toMatch('test/first.test.ts > 3 - first.test.ts - this should be skipped')
+    expect(stdout).not.toMatch('test/second.test.ts > 1 - second.test.ts - this should be skipped')
+    expect(stdout).not.toMatch('test/second.test.ts > 2 - second.test.ts - this should be skipped')
+    expect(stdout).not.toMatch('test/second.test.ts > 3 - second.test.ts - this should be skipped')
+  })
+}
diff --git a/test/bail/vitest.config.ts b/test/bail/vitest.config.ts
new file mode 100644
index 000000000000..dbce9a3386f4
--- /dev/null
+++ b/test/bail/vitest.config.ts
@@ -0,0 +1,14 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    reporters: 'verbose',
+    include: ['test/**/*.test.*'],
+
+    // For Windows CI mostly
+    testTimeout: process.env.CI ? 30_000 : 10_000,
+    chaiConfig: {
+      truncateThreshold: 999,
+    },
+  },
+})