From c9773df081a55bdee4d3e27fa18cc5c42a7d2271 Mon Sep 17 00:00:00 2001
From: AriPerkkio <ari.perkkio@gmail.com>
Date: Fri, 14 Apr 2023 16:21:28 +0300
Subject: [PATCH] feat: `--bail` option for cancelling test run

---
 packages/runner/src/types/runner.ts   |  2 +-
 packages/vitest/src/node/cli.ts       |  1 +
 packages/vitest/src/node/pools/rpc.ts |  6 ++++++
 packages/vitest/src/runtime/entry.ts  | 14 ++++++++++++++
 packages/vitest/src/types/config.ts   |  5 +++++
 packages/vitest/src/types/rpc.ts      |  2 ++
 6 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/packages/runner/src/types/runner.ts b/packages/runner/src/types/runner.ts
index e3b82241d8b4..56a257814741 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/node/cli.ts b/packages/vitest/src/node/cli.ts
index 617b20494379..9fd39a549544 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 after specified number of failures', { default: 0 })
   .help()
 
 cli
diff --git a/packages/vitest/src/node/pools/rpc.ts b/packages/vitest/src/node/pools/rpc.ts
index 1e295f9a6576..c193d19d9037 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 Array.from(ctx.state.idMap.values()).filter(t => t.result?.state === 'fail').length
+    },
   }
 }
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/types/config.ts b/packages/vitest/src/types/config.ts
index 4b6464c0c933..253d59d163c1 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. Defaults to 1.
+   */
+  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