Skip to content

Commit a3fd5f8

Browse files
authored
feat: allow custom pools (#4417)
1 parent 94f9a3c commit a3fd5f8

File tree

29 files changed

+306
-59
lines changed

29 files changed

+306
-59
lines changed

docs/.vitepress/config.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,9 +131,13 @@ export default withPwa(defineConfig({
131131
link: '/advanced/metadata',
132132
},
133133
{
134-
text: 'Extending default reporters',
134+
text: 'Extending Reporters',
135135
link: '/advanced/reporters',
136136
},
137+
{
138+
text: 'Custom Pool',
139+
link: '/advanced/pool',
140+
},
137141
],
138142
},
139143
],

docs/advanced/pool.md

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# Custom Pool
2+
3+
::: warning
4+
This is advanced API. If you are just running tests, you probably don't need this. It is primarily used by library authors.
5+
:::
6+
7+
Vitest runs tests in pools. By default, there are several pools:
8+
9+
- `threads` to run tests using `node:worker_threads` (isolation is provided with a new worker context)
10+
- `forks` to run tests using `node:child_process` (isolation is provided with a new `child_process.fork` process)
11+
- `vmThreads` to run tests using `node:worker_threads` (but isolation is provided with `vm` module instead of a new worker context)
12+
- `browser` to run tests using browser providers
13+
- `typescript` to run typechecking on tests
14+
15+
You can provide your own pool by specifying a file path:
16+
17+
```ts
18+
export default defineConfig({
19+
test: {
20+
// will run every file with a custom pool by default
21+
pool: './my-custom-pool.ts',
22+
// you can provide options using `poolOptions` object
23+
poolOptions: {
24+
myCustomPool: {
25+
customProperty: true,
26+
},
27+
},
28+
// you can also specify pool for a subset of files
29+
poolMatchGlobs: [
30+
['**/*.custom.test.ts', './my-custom-pool.ts'],
31+
],
32+
},
33+
})
34+
```
35+
36+
## API
37+
38+
The file specified in `pool` option should export a function (can be async) that accepts `Vitest` interface as its first option. This function needs to return an object matching `ProcessPool` interface:
39+
40+
```ts
41+
import { ProcessPool, WorkspaceProject } from 'vitest/node'
42+
43+
export interface ProcessPool {
44+
name: string
45+
runTests: (files: [project: WorkspaceProject, testFile: string][], invalidates?: string[]) => Promise<void>
46+
close?: () => Promise<void>
47+
}
48+
```
49+
50+
The function is called only once (unless the server config was updated), and it's generally a good idea to initialize everything you need for tests inside that function and reuse it when `runTests` is called.
51+
52+
Vitest calls `runTest` when new tests are scheduled to run. It will not call it if `files` is empty. The first argument is an array of tuples: the first element is a reference to a workspace project and the second one is an absolute path to a test file. Files are sorted using [`sequencer`](/config/#sequence.sequencer) before `runTests` is called. It's possible (but unlikely) to have the same file twice, but it will always have a different project - this is implemented via [`vitest.workspace.ts`](/guide/workspace) configuration.
53+
54+
Vitest will wait until `runTests` is executed before finishing a run (i.e., it will emit [`onFinished`](/guide/reporters) only after `runTests` is resolved).
55+
56+
If you are using a custom pool, you will have to provide test files and their results yourself - you can reference [`vitest.state`](https://github.com/vitest-dev/vitest/blob/feat/custom-pool/packages/vitest/src/node/state.ts) for that (most important are `collectFiles` and `updateTasks`). Vitest uses `startTests` function from `@vitest/runner` package to do that.
57+
58+
To communicate between different processes, you can create methods object using `createMethodsRPC` from `vitest/node`, and use any form of communication that you prefer. For example, to use websockets with `birpc` you can write something like this:
59+
60+
```ts
61+
import { createBirpc } from 'birpc'
62+
import { parse, stringify } from 'flatted'
63+
import { WorkspaceProject, createMethodsRPC } from 'vitest/node'
64+
65+
function createRpc(project: WorkspaceProject, wss: WebSocketServer) {
66+
return createBirpc(
67+
createMethodsRPC(project),
68+
{
69+
post: msg => wss.send(msg),
70+
on: fn => wss.on('message', fn),
71+
serialize: stringify,
72+
deserialize: parse,
73+
},
74+
)
75+
}
76+
```
77+
78+
To make sure every test is collected, you would call `ctx.state.collectFiles` and report it to Vitest reporters:
79+
80+
```ts
81+
async function runTests(project: WorkspaceProject, tests: string[]) {
82+
// ... running tests, put into "files" and "tasks"
83+
const methods = createMethodsRPC(project)
84+
await methods.onCollected(files)
85+
// most reporters rely on results being updated in "onTaskUpdate"
86+
await methods.onTaskUpdate(tasks)
87+
}
88+
```
89+
90+
You can see a simple example in [pool/custom-pool.ts](https://github.com/vitest-dev/vitest/blob/feat/custom-pool/test/run/pool-custom-fixtures/pool/custom-pool.ts).

docs/advanced/reporters.md

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
# Extending default reporters
1+
# Extending Reporters
22

33
You can import reporters from `vitest/reporters` and extend them to create your custom reporters.
44

5-
## Extending built-in reporters
5+
## Extending Built-in Reporters
66

77
In general, you don't need to create your reporter from scratch. `vitest` comes with several default reporting programs that you can extend.
88

@@ -56,7 +56,7 @@ export default defineConfig({
5656
})
5757
```
5858

59-
## Exported reporters
59+
## Exported Reporters
6060

6161
`vitest` comes with a few [built-in reporters](/guide/reporters) that you can use out of the box.
6262

packages/runner/src/suite.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
154154
shuffle,
155155
tasks: [],
156156
meta: Object.create(null),
157+
projectName: '',
157158
}
158159

159160
setHooks(suite, createSuiteHooks())

packages/runner/src/types/tasks.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ export interface Suite extends TaskBase {
5252
type: 'suite'
5353
tasks: Task[]
5454
filepath?: string
55-
projectName?: string
55+
projectName: string
5656
}
5757

5858
export interface File extends Suite {

packages/vitest/src/node/config.ts

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@ import type { ApiConfig, ResolvedConfig, UserConfig, VitestRunMode } from '../ty
66
import { defaultBrowserPort, defaultPort } from '../constants'
77
import { benchmarkConfigDefaults, configDefaults } from '../defaults'
88
import { isCI, stdProvider, toArray } from '../utils'
9+
import type { BuiltinPool } from '../types/pool-options'
910
import { VitestCache } from './cache'
1011
import { BaseSequencer } from './sequencers/BaseSequencer'
1112
import { RandomSequencer } from './sequencers/RandomSequencer'
1213
import type { BenchmarkBuiltinReporters } from './reporters'
14+
import { builtinPools } from './pool'
1315

1416
const extraInlineDeps = [
1517
/^(?!.*(?:node_modules)).*\.mjs$/,
@@ -222,6 +224,8 @@ export function resolveConfig(
222224
if (options.resolveSnapshotPath)
223225
delete (resolved as UserConfig).resolveSnapshotPath
224226

227+
resolved.pool ??= 'threads'
228+
225229
if (process.env.VITEST_MAX_THREADS) {
226230
resolved.poolOptions = {
227231
...resolved.poolOptions,
@@ -270,6 +274,22 @@ export function resolveConfig(
270274
}
271275
}
272276

277+
if (!builtinPools.includes(resolved.pool as BuiltinPool)) {
278+
resolved.pool = normalize(
279+
resolveModule(resolved.pool, { paths: [resolved.root] })
280+
?? resolve(resolved.root, resolved.pool),
281+
)
282+
}
283+
resolved.poolMatchGlobs = (resolved.poolMatchGlobs || []).map(([glob, pool]) => {
284+
if (!builtinPools.includes(pool as BuiltinPool)) {
285+
pool = normalize(
286+
resolveModule(pool, { paths: [resolved.root] })
287+
?? resolve(resolved.root, pool),
288+
)
289+
}
290+
return [glob, pool]
291+
})
292+
273293
if (mode === 'benchmark') {
274294
resolved.benchmark = {
275295
...benchmarkConfigDefaults,

packages/vitest/src/node/core.ts

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ export class Vitest {
8080
this.unregisterWatcher?.()
8181
clearTimeout(this._rerunTimer)
8282
this.restartsCount += 1
83-
this.pool?.close()
83+
this.pool?.close?.()
8484
this.pool = undefined
8585
this.coverageProvider = undefined
8686
this.runningPromise = undefined
@@ -761,8 +761,12 @@ export class Vitest {
761761
if (!this.projects.includes(this.coreWorkspaceProject))
762762
closePromises.push(this.coreWorkspaceProject.close().then(() => this.server = undefined as any))
763763

764-
if (this.pool)
765-
closePromises.push(this.pool.close().then(() => this.pool = undefined))
764+
if (this.pool) {
765+
closePromises.push((async () => {
766+
await this.pool?.close?.()
767+
this.pool = undefined
768+
})())
769+
}
766770

767771
closePromises.push(...this._onClose.map(fn => fn()))
768772

packages/vitest/src/node/index.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@ export { createVitest } from './create'
44
export { VitestPlugin } from './plugins'
55
export { startVitest } from './cli-api'
66
export { registerConsoleShortcuts } from './stdin'
7-
export type { WorkspaceSpec } from './pool'
87
export type { GlobalSetupContext } from './globalSetup'
8+
export type { WorkspaceSpec, ProcessPool } from './pool'
9+
export { createMethodsRPC } from './pools/rpc'
910

1011
export type { TestSequencer, TestSequencerConstructor } from './sequencers/types'
1112
export { BaseSequencer } from './sequencers/BaseSequencer'

packages/vitest/src/node/pool.ts

Lines changed: 51 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import mm from 'micromatch'
2-
import type { Pool } from '../types'
2+
import type { Awaitable } from '@vitest/utils'
3+
import type { BuiltinPool, Pool } from '../types/pool-options'
34
import type { Vitest } from './core'
45
import { createChildProcessPool } from './pools/child'
56
import { createThreadsPool } from './pools/threads'
@@ -9,11 +10,12 @@ import type { WorkspaceProject } from './workspace'
910
import { createTypecheckPool } from './pools/typecheck'
1011

1112
export type WorkspaceSpec = [project: WorkspaceProject, testFile: string]
12-
export type RunWithFiles = (files: WorkspaceSpec[], invalidates?: string[]) => Promise<void>
13+
export type RunWithFiles = (files: WorkspaceSpec[], invalidates?: string[]) => Awaitable<void>
1314

1415
export interface ProcessPool {
16+
name: string
1517
runTests: RunWithFiles
16-
close: () => Promise<void>
18+
close?: () => Awaitable<void>
1719
}
1820

1921
export interface PoolProcessOptions {
@@ -24,6 +26,8 @@ export interface PoolProcessOptions {
2426
env: Record<string, string>
2527
}
2628

29+
export const builtinPools: BuiltinPool[] = ['forks', 'threads', 'browser', 'vmThreads', 'typescript']
30+
2731
export function createPool(ctx: Vitest): ProcessPool {
2832
const pools: Record<Pool, ProcessPool | null> = {
2933
forks: null,
@@ -48,7 +52,7 @@ export function createPool(ctx: Vitest): ProcessPool {
4852
}
4953

5054
function getPoolName([project, file]: WorkspaceSpec) {
51-
for (const [glob, pool] of project.config.poolMatchGlobs || []) {
55+
for (const [glob, pool] of project.config.poolMatchGlobs) {
5256
if ((pool as Pool) === 'browser')
5357
throw new Error('Since Vitest 0.31.0 "browser" pool is not supported in "poolMatchGlobs". You can create a workspace to run some of your tests in browser in parallel. Read more: https://vitest.dev/guide/workspace')
5458
if (mm.isMatch(file, glob, { cwd: project.config.root }))
@@ -82,6 +86,22 @@ export function createPool(ctx: Vitest): ProcessPool {
8286
},
8387
}
8488

89+
const customPools = new Map<string, ProcessPool>()
90+
async function resolveCustomPool(filepath: string) {
91+
if (customPools.has(filepath))
92+
return customPools.get(filepath)!
93+
const pool = await ctx.runner.executeId(filepath)
94+
if (typeof pool.default !== 'function')
95+
throw new Error(`Custom pool "${filepath}" must export a function as default export`)
96+
const poolInstance = await pool.default(ctx, options)
97+
if (typeof poolInstance?.name !== 'string')
98+
throw new Error(`Custom pool "${filepath}" should return an object with "name" property`)
99+
if (typeof poolInstance?.runTests !== 'function')
100+
throw new Error(`Custom pool "${filepath}" should return an object with "runTests" method`)
101+
customPools.set(filepath, poolInstance)
102+
return poolInstance as ProcessPool
103+
}
104+
85105
const filesByPool: Record<Pool, WorkspaceSpec[]> = {
86106
forks: [],
87107
threads: [],
@@ -92,46 +112,63 @@ export function createPool(ctx: Vitest): ProcessPool {
92112

93113
for (const spec of files) {
94114
const pool = getPoolName(spec)
95-
if (!(pool in filesByPool))
96-
throw new Error(`Unknown pool name "${pool}" for ${spec[1]}. Available pools: ${Object.keys(filesByPool).join(', ')}`)
115+
filesByPool[pool] ??= []
97116
filesByPool[pool].push(spec)
98117
}
99118

100-
await Promise.all(Object.entries(filesByPool).map((entry) => {
119+
const Sequencer = ctx.config.sequence.sequencer
120+
const sequencer = new Sequencer(ctx)
121+
122+
async function sortSpecs(specs: WorkspaceSpec[]) {
123+
if (ctx.config.shard)
124+
specs = await sequencer.shard(specs)
125+
return sequencer.sort(specs)
126+
}
127+
128+
await Promise.all(Object.entries(filesByPool).map(async (entry) => {
101129
const [pool, files] = entry as [Pool, WorkspaceSpec[]]
102130

103131
if (!files.length)
104132
return null
105133

134+
const specs = await sortSpecs(files)
135+
106136
if (pool === 'browser') {
107137
pools.browser ??= createBrowserPool(ctx)
108-
return pools.browser.runTests(files, invalidate)
138+
return pools.browser.runTests(specs, invalidate)
109139
}
110140

111141
if (pool === 'vmThreads') {
112142
pools.vmThreads ??= createVmThreadsPool(ctx, options)
113-
return pools.vmThreads.runTests(files, invalidate)
143+
return pools.vmThreads.runTests(specs, invalidate)
114144
}
115145

116146
if (pool === 'threads') {
117147
pools.threads ??= createThreadsPool(ctx, options)
118-
return pools.threads.runTests(files, invalidate)
148+
return pools.threads.runTests(specs, invalidate)
119149
}
120150

121151
if (pool === 'typescript') {
122152
pools.typescript ??= createTypecheckPool(ctx)
123-
return pools.typescript.runTests(files)
153+
return pools.typescript.runTests(specs)
154+
}
155+
156+
if (pool === 'forks') {
157+
pools.forks ??= createChildProcessPool(ctx, options)
158+
return pools.forks.runTests(specs, invalidate)
124159
}
125160

126-
pools.forks ??= createChildProcessPool(ctx, options)
127-
return pools.forks.runTests(files, invalidate)
161+
const poolHandler = await resolveCustomPool(pool)
162+
pools[poolHandler.name] ??= poolHandler
163+
return poolHandler.runTests(specs, invalidate)
128164
}))
129165
}
130166

131167
return {
168+
name: 'default',
132169
runTests,
133170
async close() {
134-
await Promise.all(Object.values(pools).map(p => p?.close()))
171+
await Promise.all(Object.values(pools).map(p => p?.close?.()))
135172
},
136173
}
137174
}

packages/vitest/src/node/pools/browser.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
4444
if (project.config.browser.isolate) {
4545
for (const path of paths) {
4646
if (isCancelled) {
47-
ctx.state.cancelFiles(files.slice(paths.indexOf(path)), ctx.config.root)
47+
ctx.state.cancelFiles(files.slice(paths.indexOf(path)), ctx.config.root, project.getName())
4848
break
4949
}
5050

@@ -77,6 +77,7 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
7777
}
7878

7979
return {
80+
name: 'browser',
8081
async close() {
8182
ctx.state.browserTestPromises.clear()
8283
await Promise.all([...providers].map(provider => provider.close()))

0 commit comments

Comments
 (0)