From e189122dee8bb6f3c7339329c97b5912e3bdcec4 Mon Sep 17 00:00:00 2001 From: AriPerkkio Date: Tue, 11 Apr 2023 15:03:26 +0300 Subject: [PATCH] feat(watch): test run cancelling --- docs/advanced/runner.md | 9 ++++- packages/browser/src/client/main.ts | 15 +++++++- packages/runner/src/types/runner.ts | 7 ++++ packages/vitest/src/api/setup.ts | 4 ++- packages/vitest/src/api/types.ts | 1 + packages/vitest/src/node/core.ts | 14 +++++++- packages/vitest/src/node/pools/child.ts | 7 ++-- packages/vitest/src/node/pools/threads.ts | 40 ++++++++++++++++----- packages/vitest/src/node/reporters/base.ts | 5 +++ packages/vitest/src/node/stdin.ts | 14 ++++---- packages/vitest/src/runtime/child.ts | 16 +++++++-- packages/vitest/src/runtime/entry.ts | 1 + packages/vitest/src/runtime/runners/test.ts | 13 +++++++ packages/vitest/src/runtime/worker.ts | 16 +++++++-- packages/vitest/src/types/rpc.ts | 5 +++ packages/vitest/src/types/worker.ts | 1 + packages/ws-client/src/index.ts | 3 ++ 17 files changed, 145 insertions(+), 26 deletions(-) diff --git a/docs/advanced/runner.md b/docs/advanced/runner.md index a03b6d74ea28..dc4ae5ea257e 100644 --- a/docs/advanced/runner.md +++ b/docs/advanced/runner.md @@ -17,6 +17,13 @@ export interface VitestRunner { */ onCollected?(files: File[]): unknown + /** + * Called when test runner should cancel next test runs. + * Runner should listen for this method and mark tests and suites as skipped in + * "onBeforeRunSuite" and "onBeforeRunTest" when called. + */ + onCancel?(): unknown + /** * Called before running a single test. Doesn't have "result" yet. */ @@ -86,7 +93,7 @@ export interface VitestRunner { When initiating this class, Vitest passes down Vitest config, - you should expose it as a `config` property. ::: warning -Vitest also injects an instance of `ViteNodeRunner` as `__vitest_executor` property. You can use it to process files in `importFile` method (this is default behavior of `TestRunner`` and `BenchmarkRunner`). +Vitest also injects an instance of `ViteNodeRunner` as `__vitest_executor` property. You can use it to process files in `importFile` method (this is default behavior of `TestRunner` and `BenchmarkRunner`). `ViteNodeRunner` exposes `executeId` method, which is used to import test files in a Vite-friendly environment. Meaning, it will resolve imports and transform file content at runtime so that Node can understand it. ::: diff --git a/packages/browser/src/client/main.ts b/packages/browser/src/client/main.ts index 50a636cb7ca5..78693dca9b76 100644 --- a/packages/browser/src/client/main.ts +++ b/packages/browser/src/client/main.ts @@ -30,7 +30,16 @@ function getQueryPaths() { return url.searchParams.getAll('path') } -export const client = createClient(ENTRY_URL) +let setCancel = () => {} +const onCancel = new Promise((resolve) => { + setCancel = resolve +}) + +export const client = createClient(ENTRY_URL, { + handlers: { + onCancel: setCancel, + }, +}) const ws = client.ws @@ -103,6 +112,10 @@ async function runTests(paths: string[], config: ResolvedConfig) { runner = new BrowserRunner({ config, browserHashMap }) } + onCancel.then(() => { + runner?.onCancel?.() + }) + if (!config.snapshotOptions.snapshotEnvironment) config.snapshotOptions.snapshotEnvironment = new BrowserSnapshotEnvironment() diff --git a/packages/runner/src/types/runner.ts b/packages/runner/src/types/runner.ts index ef30599a864a..2f67d72be061 100644 --- a/packages/runner/src/types/runner.ts +++ b/packages/runner/src/types/runner.ts @@ -34,6 +34,13 @@ export interface VitestRunner { */ onCollected?(files: File[]): unknown + /** + * Called when test runner should cancel next test runs. + * Runner should listen for this method and mark tests and suites as skipped in + * "onBeforeRunSuite" and "onBeforeRunTest" when called. + */ + onCancel?(): unknown + /** * Called before running a single test. Doesn't have "result" yet. */ diff --git a/packages/vitest/src/api/setup.ts b/packages/vitest/src/api/setup.ts index afff7a4bc439..032625a5855d 100644 --- a/packages/vitest/src/api/setup.ts +++ b/packages/vitest/src/api/setup.ts @@ -112,12 +112,14 @@ export function setup(vitestOrWorkspace: Vitest | WorkspaceProject, server?: Vit { post: msg => ws.send(msg), on: fn => ws.on('message', fn), - eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected'], + eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onCancel'], serialize: stringify, deserialize: parse, }, ) + ctx.onCancel(() => rpc.onCancel()) + clients.set(ws, rpc) ws.on('close', () => { diff --git a/packages/vitest/src/api/types.ts b/packages/vitest/src/api/types.ts index feb11a392707..699b4ad21523 100644 --- a/packages/vitest/src/api/types.ts +++ b/packages/vitest/src/api/types.ts @@ -28,4 +28,5 @@ export interface WebSocketHandlers { } export interface WebSocketEvents extends Pick { + onCancel(): void } diff --git a/packages/vitest/src/node/core.ts b/packages/vitest/src/node/core.ts index add6e417f39e..2e61130324be 100644 --- a/packages/vitest/src/node/core.ts +++ b/packages/vitest/src/node/core.ts @@ -46,6 +46,7 @@ export class Vitest { filenamePattern?: string runningPromise?: Promise closingPromise?: Promise + isCancelling = false isFirstRun = true restartsCount = 0 @@ -64,6 +65,7 @@ export class Vitest { private _onRestartListeners: OnServerRestartHandler[] = [] private _onSetServer: OnServerRestartHandler[] = [] + private _onCancelListeners: (() => Promise | void)[] = [] async setServer(options: UserConfig, server: ViteDevServer, cliOptions: UserConfig) { this.unregisterWatcher?.() @@ -394,13 +396,14 @@ export class Vitest { async runFiles(paths: WorkspaceSpec[]) { const filepaths = paths.map(([, file]) => file) - this.state.collectPaths(filepaths) await this.report('onPathsCollected', filepaths) // previous run await this.runningPromise + this._onCancelListeners = [] + this.isCancelling = false // schedule the new run this.runningPromise = (async () => { @@ -436,6 +439,11 @@ export class Vitest { return await this.runningPromise } + async cancelCurrentRun() { + this.isCancelling = true + await Promise.all(this._onCancelListeners.splice(0).map(listener => listener())) + } + async rerunFiles(files: string[] = this.state.getFilepaths(), trigger?: string) { if (this.filenamePattern) { const filteredFiles = await this.globTestFiles([this.filenamePattern]) @@ -765,4 +773,8 @@ export class Vitest { onAfterSetServer(fn: OnServerRestartHandler) { this._onSetServer.push(fn) } + + onCancel(fn: () => void) { + this._onCancelListeners.push(fn) + } } diff --git a/packages/vitest/src/node/pools/child.ts b/packages/vitest/src/node/pools/child.ts index 33fac5f414c3..8c48333597e3 100644 --- a/packages/vitest/src/node/pools/child.ts +++ b/packages/vitest/src/node/pools/child.ts @@ -4,7 +4,7 @@ import { fork } from 'node:child_process' import { fileURLToPath, pathToFileURL } from 'node:url' import { createBirpc } from 'birpc' import { resolve } from 'pathe' -import type { ContextTestEnvironment, ResolvedConfig, RuntimeRPC, Vitest } from '../../types' +import type { ContextTestEnvironment, ResolvedConfig, RunnerRPC, RuntimeRPC, Vitest } from '../../types' import type { ChildContext } from '../../types/child' import type { PoolProcessOptions, ProcessPool, WorkspaceSpec } from '../pool' import { distDir } from '../../paths' @@ -16,9 +16,10 @@ import { createMethodsRPC } from './rpc' const childPath = fileURLToPath(pathToFileURL(resolve(distDir, './child.js')).href) function setupChildProcessChannel(project: WorkspaceProject, fork: ChildProcess): void { - createBirpc<{}, RuntimeRPC>( + const rpc = createBirpc( createMethodsRPC(project), { + eventNames: ['onCancel'], serialize: v8.serialize, deserialize: v => v8.deserialize(Buffer.from(v)), post(v) { @@ -29,6 +30,8 @@ function setupChildProcessChannel(project: WorkspaceProject, fork: ChildProcess) }, }, ) + + project.ctx.onCancel(() => rpc.onCancel()) } function stringifyRegex(input: RegExp | string): string { diff --git a/packages/vitest/src/node/pools/threads.ts b/packages/vitest/src/node/pools/threads.ts index 4961b15c3070..b31a2aa9abc5 100644 --- a/packages/vitest/src/node/pools/threads.ts +++ b/packages/vitest/src/node/pools/threads.ts @@ -1,12 +1,13 @@ import { MessageChannel } from 'node:worker_threads' import { cpus } from 'node:os' import { pathToFileURL } from 'node:url' +import EventEmitter from 'node:events' import { createBirpc } from 'birpc' import { resolve } from 'pathe' import type { Options as TinypoolOptions } from 'tinypool' import Tinypool from 'tinypool' import { distDir } from '../../paths' -import type { ContextTestEnvironment, ResolvedConfig, RuntimeRPC, Vitest, WorkerContext } from '../../types' +import type { ContextTestEnvironment, ResolvedConfig, RunnerRPC, RuntimeRPC, Vitest, WorkerContext } from '../../types' import type { PoolProcessOptions, ProcessPool, RunWithFiles } from '../pool' import { envsOrder, groupFilesByEnv } from '../../utils/test-helpers' import { groupBy } from '../../utils/base' @@ -20,9 +21,10 @@ function createWorkerChannel(project: WorkspaceProject) { const port = channel.port2 const workerPort = channel.port1 - createBirpc<{}, RuntimeRPC>( + const rpc = createBirpc( createMethodsRPC(project), { + eventNames: ['onCancel'], post(v) { port.postMessage(v) }, @@ -32,6 +34,8 @@ function createWorkerChannel(project: WorkspaceProject) { }, ) + project.ctx.onCancel(() => rpc.onCancel()) + return { workerPort, port } } @@ -74,7 +78,7 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt const runWithFiles = (name: string): RunWithFiles => { let id = 0 - async function runFiles(project: WorkspaceProject, config: ResolvedConfig, files: string[], environment: ContextTestEnvironment, invalidates: string[] = []) { + async function runFiles(project: WorkspaceProject, config: ResolvedConfig, files: string[], environment: ContextTestEnvironment, invalidates: string[] = [], signal: EventEmitter) { ctx.state.clearFiles(project, files) const { workerPort, port } = createWorkerChannel(project) const workerId = ++id @@ -87,14 +91,28 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt workerId, } try { - await pool.run(data, { transferList: [workerPort], name }) + await pool.run(data, { transferList: [workerPort], name, signal }) } catch (error) { // Worker got stuck and won't terminate - this may cause process to hang - if (error instanceof Error && /Failed to terminate worker/.test(error.message)) + if (error instanceof Error && /Failed to terminate worker/.test(error.message)) { ctx.state.addProcessTimeoutCause(`Failed to terminate worker while running ${files.join(', ')}.`) - else + } + else if (ctx.isCancelling && error instanceof Error && /The task has been aborted/.test(error.message)) { + // Intentionally cancelled + // TODO: This could be a "ctx.state.cancelFiles(files: string[])" instead + ctx.state.collectFiles(files.map(filepath => ({ + filepath, + id: filepath, + mode: 'skip', + name: filepath, + tasks: [], + type: 'suite', + }))) + } + else { throw error + } } finally { port.close() @@ -123,6 +141,12 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt workspaceMap.set(file, workspaceFiles) } + // TODO: Check and make sure Tinypool cancels tasks gracefully instead of forcefully terminating them. + // If not, remove whole signal and rely on runner level cancellation instead. + const signal = new EventEmitter() + signal.setMaxListeners(ctx.state.getFiles().length) + ctx.onCancel(() => signal.emit('abort')) + // it's possible that project defines a file that is also defined by another project const { shard } = ctx.config @@ -138,7 +162,7 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt const filesByEnv = await groupFilesByEnv(multipleThreads) const promises = Object.values(filesByEnv).flat() const results = await Promise.allSettled(promises - .map(({ file, environment, project }) => runFiles(project, getConfig(project), [file], environment, invalidates))) + .map(({ file, environment, project }) => runFiles(project, getConfig(project), [file], environment, invalidates, signal))) const errors = results.filter((r): r is PromiseRejectedResult => r.status === 'rejected').map(r => r.reason) if (errors.length > 0) @@ -162,7 +186,7 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt const promises = Object.values(filesByOptions).map(async (files) => { const filenames = files.map(f => f.file) - await runFiles(files[0].project, getConfig(files[0].project), filenames, files[0].environment, invalidates) + await runFiles(files[0].project, getConfig(files[0].project), filenames, files[0].environment, invalidates, signal) }) await Promise.all(promises) diff --git a/packages/vitest/src/node/reporters/base.ts b/packages/vitest/src/node/reporters/base.ts index 15087eec1229..484840208bec 100644 --- a/packages/vitest/src/node/reporters/base.ts +++ b/packages/vitest/src/node/reporters/base.ts @@ -13,6 +13,7 @@ const HELP_QUITE = `${c.dim('press ')}${c.bold('q')}${c.dim(' to quit')}` const WAIT_FOR_CHANGE_PASS = `\n${c.bold(c.inverse(c.green(' PASS ')))}${c.green(' Waiting for file changes...')}` const WAIT_FOR_CHANGE_FAIL = `\n${c.bold(c.inverse(c.red(' FAIL ')))}${c.red(' Tests failed. Watching for file changes...')}` +const WAIT_FOR_CHANGE_CANCELLED = `\n${c.bold(c.inverse(c.red(' CANCELLED ')))}${c.red(' Test run cancelled. Watching for file changes...')}` const LAST_RUN_LOG_TIMEOUT = 1_500 @@ -102,8 +103,12 @@ export abstract class BaseReporter implements Reporter { const failed = errors.length > 0 || hasFailed(files) const failedSnap = hasFailedSnapshot(files) + const cancelled = this.ctx.isCancelling + if (failed) this.ctx.logger.log(WAIT_FOR_CHANGE_FAIL) + else if (cancelled) + this.ctx.logger.log(WAIT_FOR_CHANGE_CANCELLED) else this.ctx.logger.log(WAIT_FOR_CHANGE_PASS) diff --git a/packages/vitest/src/node/stdin.ts b/packages/vitest/src/node/stdin.ts index dfcc5d4824e4..692e925771e2 100644 --- a/packages/vitest/src/node/stdin.ts +++ b/packages/vitest/src/node/stdin.ts @@ -37,12 +37,14 @@ export function registerConsoleShortcuts(ctx: Vitest) { return } - // is running, ignore keypress - if (ctx.runningPromise) - return - const name = key?.name + if (ctx.runningPromise) { + if (['space', 'c'].includes(name) || keys.map(key => key[0]).includes(name)) + await ctx.cancelCurrentRun() + return + } + // quit if (name === 'q') return ctx.exit(true) @@ -83,8 +85,8 @@ export function registerConsoleShortcuts(ctx: Vitest) { message: 'Input test name pattern (RegExp)', initial: ctx.configOverride.testNamePattern?.source || '', }]) - await ctx.changeNamePattern(filter.trim(), undefined, 'change pattern') on() + await ctx.changeNamePattern(filter.trim(), undefined, 'change pattern') } async function inputFilePattern() { @@ -96,8 +98,8 @@ export function registerConsoleShortcuts(ctx: Vitest) { initial: latestFilename, }]) latestFilename = filter.trim() - await ctx.changeFilenamePattern(filter.trim()) on() + await ctx.changeFilenamePattern(filter.trim()) } let rl: readline.Interface | undefined diff --git a/packages/vitest/src/runtime/child.ts b/packages/vitest/src/runtime/child.ts index 29a6d690333b..228776e5311b 100644 --- a/packages/vitest/src/runtime/child.ts +++ b/packages/vitest/src/runtime/child.ts @@ -2,7 +2,7 @@ import v8 from 'node:v8' import { createBirpc } from 'birpc' import { parseRegexp } from '@vitest/utils' import type { ResolvedConfig } from '../types' -import type { RuntimeRPC } from '../types/rpc' +import type { RunnerRPC, RuntimeRPC } from '../types/rpc' import type { ChildContext } from '../types/child' import { mockMap, moduleCache, startViteNode } from './execute' import { rpcDone } from './rpc' @@ -14,6 +14,11 @@ function init(ctx: ChildContext) { process.env.VITEST_WORKER_ID = '1' process.env.VITEST_POOL_ID = '1' + let setCancel = () => {} + const onCancel = new Promise((resolve) => { + setCancel = resolve + }) + // @ts-expect-error untyped global globalThis.__vitest_environment__ = config.environment // @ts-expect-error I know what I am doing :P @@ -22,12 +27,17 @@ function init(ctx: ChildContext) { moduleCache, config, mockMap, + onCancel, durations: { environment: 0, prepare: performance.now(), }, - rpc: createBirpc( - {}, + rpc: createBirpc( + { + onCancel() { + setCancel() + }, + }, { eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'], serialize: v8.serialize, diff --git a/packages/vitest/src/runtime/entry.ts b/packages/vitest/src/runtime/entry.ts index e16f75285214..5134f67405ec 100644 --- a/packages/vitest/src/runtime/entry.ts +++ b/packages/vitest/src/runtime/entry.ts @@ -84,6 +84,7 @@ export async function run(files: string[], config: ResolvedConfig, environment: setupChaiConfig(config.chaiConfig) const runner = await getTestRunner(config, executor) + workerState.onCancel.then(() => runner.onCancel?.()) workerState.durations.prepare = performance.now() - workerState.durations.prepare diff --git a/packages/vitest/src/runtime/runners/test.ts b/packages/vitest/src/runtime/runners/test.ts index 1abcfce709ce..8b51519192bc 100644 --- a/packages/vitest/src/runtime/runners/test.ts +++ b/packages/vitest/src/runtime/runners/test.ts @@ -12,6 +12,7 @@ export class VitestTestRunner implements VitestRunner { private snapshotClient = getSnapshotClient() private workerState = getWorkerState() private __vitest_executor!: VitestExecutor + private cancelRun = false constructor(public config: ResolvedConfig) {} @@ -45,9 +46,16 @@ export class VitestTestRunner implements VitestRunner { this.workerState.current = undefined } + onCancel() { + this.cancelRun = true + } + async onBeforeRunTest(test: Test) { const name = getNames(test).slice(1).join(' > ') + if (this.cancelRun) + test.mode = 'skip' + if (test.mode !== 'run') { this.snapshotClient.skipTestSnapshots(name) return @@ -59,6 +67,11 @@ export class VitestTestRunner implements VitestRunner { this.workerState.current = test } + onBeforeRunSuite(suite: Suite) { + if (this.cancelRun) + suite.mode = 'skip' + } + onBeforeTryTest(test: Test) { setState({ assertionCalls: 0, diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts index 586d9390067c..4f282db1d6d3 100644 --- a/packages/vitest/src/runtime/worker.ts +++ b/packages/vitest/src/runtime/worker.ts @@ -1,6 +1,6 @@ import { createBirpc } from 'birpc' import { workerId as poolId } from 'tinypool' -import type { RuntimeRPC, WorkerContext } from '../types' +import type { RunnerRPC, RuntimeRPC, WorkerContext } from '../types' import { getWorkerState } from '../utils/global' import { mockMap, moduleCache, startViteNode } from './execute' import { setupInspect } from './inspector' @@ -16,6 +16,11 @@ function init(ctx: WorkerContext) { process.env.VITEST_WORKER_ID = String(workerId) process.env.VITEST_POOL_ID = String(poolId) + let setCancel = () => {} + const onCancel = new Promise((resolve) => { + setCancel = resolve + }) + // @ts-expect-error untyped global globalThis.__vitest_environment__ = config.environment // @ts-expect-error I know what I am doing :P @@ -24,12 +29,17 @@ function init(ctx: WorkerContext) { moduleCache, config, mockMap, + onCancel, durations: { environment: 0, prepare: performance.now(), }, - rpc: createBirpc( - {}, + rpc: createBirpc( + { + onCancel() { + setCancel() + }, + }, { eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'], post(v) { port.postMessage(v) }, diff --git a/packages/vitest/src/types/rpc.ts b/packages/vitest/src/types/rpc.ts index eb10ba97ed35..2281131b9462 100644 --- a/packages/vitest/src/types/rpc.ts +++ b/packages/vitest/src/types/rpc.ts @@ -24,6 +24,11 @@ export interface RuntimeRPC { resolveSnapshotPath: (testPath: string) => string } +export interface RunnerRPC { + // TODO: This could be "(reason: 'INPUT' | 'BAIL') => void" instead + onCancel: () => void +} + export interface ContextTestEnvironment { name: VitestEnvironment options: EnvironmentOptions | null diff --git a/packages/vitest/src/types/worker.ts b/packages/vitest/src/types/worker.ts index cc0face52dbe..be5e07267571 100644 --- a/packages/vitest/src/types/worker.ts +++ b/packages/vitest/src/types/worker.ts @@ -24,6 +24,7 @@ export interface WorkerGlobalState { current?: Test filepath?: string environmentTeardownRun?: boolean + onCancel: Promise moduleCache: ModuleCacheMap mockMap: MockMap durations: { diff --git a/packages/ws-client/src/index.ts b/packages/ws-client/src/index.ts index d016e12c8cee..62cb8c5cb47d 100644 --- a/packages/ws-client/src/index.ts +++ b/packages/ws-client/src/index.ts @@ -66,6 +66,9 @@ export function createClient(url: string, options: VitestClientOptions = {}) { onFinished(files) { handlers.onFinished?.(files) }, + onCancel() { + handlers.onCancel?.() + }, } const birpcHandlers: BirpcOptions = {