feat: add an option to control Vitest pool with filepath (#3029)

Author: sheremet-va

docs/config/index.md

@@ -381,6 +381,31 @@ export default defineConfig({
 })
 ```
 
+### poolMatchGlobs
+
+- **Type:** `[string, 'threads' | 'child_process'][]`
+- **Default:** `[]`
+- **Version:** Since Vitest 0.29.4
+
+Automatically assign pool in which tests will run based on globs. The first match will be used.
+
+For example:
+
+```ts
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    poolMatchGlobs: [
+      // all tests in "worker-specific" directory will run inside a worker as if you enabled `--threads` for them,
+      ['**/tests/worker-specific/**', 'threads'],
+      // all other tests will run based on "threads" option, if you didn't specify other globs
+      // ...
+    ]
+  }
+})
+```
+
 ### update
 
 - **Type:** `boolean`

packages/vitest/src/node/pool.ts

@@ -1,6 +1,8 @@
 import { pathToFileURL } from 'node:url'
+import mm from 'micromatch'
 import { resolve } from 'pathe'
 import { distDir, rootDir } from '../constants'
+import type { VitestPool } from '../types'
 import type { Vitest } from './core'
 import { createChildProcessPool } from './pools/child'
 import { createThreadsPool } from './pools/threads'
@@ -21,38 +23,92 @@ const loaderPath = pathToFileURL(resolve(distDir, './loader.js')).href
 const suppressLoaderWarningsPath = resolve(rootDir, './suppress-warnings.cjs')
 
 export function createPool(ctx: Vitest): ProcessPool {
-  const conditions = ctx.server.config.resolve.conditions?.flatMap(c => ['--conditions', c]) || []
-
-  // Instead of passing whole process.execArgv to the workers, pick allowed options.
-  // Some options may crash worker, e.g. --prof, --title. nodejs/node#41103
-  const execArgv = process.execArgv.filter(execArg =>
-    execArg.startsWith('--cpu-prof') || execArg.startsWith('--heap-prof'),
-  )
-
-  const options: PoolProcessOptions = {
-    execArgv: ctx.config.deps.registerNodeLoader
-      ? [
-          ...execArgv,
-          '--require',
-          suppressLoaderWarningsPath,
-          '--experimental-loader',
-          loaderPath,
-        ]
-      : [
-          ...execArgv,
-          ...conditions,
-        ],
-    env: {
-      TEST: 'true',
-      VITEST: 'true',
-      NODE_ENV: ctx.config.mode || 'test',
-      VITEST_MODE: ctx.config.watch ? 'WATCH' : 'RUN',
-      ...process.env,
-      ...ctx.config.env,
-    },
+  const pools: Record<VitestPool, ProcessPool | null> = {
+    child_process: null,
+    threads: null,
+  }
+
+  function getDefaultPoolName() {
+    if (ctx.config.threads)
+      return 'threads'
+    return 'child_process'
+  }
+
+  function getPoolName(file: string) {
+    for (const [glob, pool] of ctx.config.poolMatchGlobs || []) {
+      if (mm.isMatch(file, glob, { cwd: ctx.server.config.root }))
+        return pool
+    }
+    return getDefaultPoolName()
+  }
+
+  async function runTests(files: string[], invalidate?: string[]) {
+    const conditions = ctx.server.config.resolve.conditions?.flatMap(c => ['--conditions', c]) || []
+
+    // Instead of passing whole process.execArgv to the workers, pick allowed options.
+    // Some options may crash worker, e.g. --prof, --title. nodejs/node#41103
+    const execArgv = process.execArgv.filter(execArg =>
+      execArg.startsWith('--cpu-prof') || execArg.startsWith('--heap-prof'),
+    )
+
+    const options: PoolProcessOptions = {
+      execArgv: ctx.config.deps.registerNodeLoader
+        ? [
+            ...execArgv,
+            '--require',
+            suppressLoaderWarningsPath,
+            '--experimental-loader',
+            loaderPath,
+          ]
+        : [
+            ...execArgv,
+            ...conditions,
+          ],
+      env: {
+        TEST: 'true',
+        VITEST: 'true',
+        NODE_ENV: ctx.config.mode || 'test',
+        VITEST_MODE: ctx.config.watch ? 'WATCH' : 'RUN',
+        ...process.env,
+        ...ctx.config.env,
+      },
+    }
+
+    const filesByPool = {
+      child_process: [] as string[],
+      threads: [] as string[],
+      browser: [] as string[],
+    }
+
+    if (!ctx.config.poolMatchGlobs) {
+      const name = getDefaultPoolName()
+      filesByPool[name] = files
+    }
+    else {
+      for (const file of files) {
+        const pool = getPoolName(file)
+        filesByPool[pool].push(file)
+      }
+    }
+
+    await Promise.all(Object.entries(filesByPool).map(([pool, files]) => {
+      if (!files.length)
+        return null
+
+      if (pool === 'threads') {
+        pools.threads ??= createThreadsPool(ctx, options)
+        return pools.threads.runTests(files, invalidate)
+      }
+
+      pools.child_process ??= createChildProcessPool(ctx, options)
+      return pools.child_process.runTests(files, invalidate)
+    }))
   }
 
-  if (!ctx.config.threads)
-    return createChildProcessPool(ctx, options)
-  return createThreadsPool(ctx, options)
+  return {
+    runTests,
+    async close() {
+      await Promise.all(Object.values(pools).map(p => p?.close()))
+    },
+  }
 }

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

@@ -15,7 +15,7 @@ import { createMethodsRPC } from './rpc'
 
 const childPath = fileURLToPath(pathToFileURL(resolve(distDir, './child.js')).href)
 
-function setupChildProcessChannel(ctx: Vitest, fork: ChildProcess) {
+function setupChildProcessChannel(ctx: Vitest, fork: ChildProcess): void {
   createBirpc<{}, RuntimeRPC>(
     createMethodsRPC(ctx),
     {
@@ -31,16 +31,16 @@ function setupChildProcessChannel(ctx: Vitest, fork: ChildProcess) {
   )
 }
 
-function stringifyRegex(input: RegExp | string): any {
+function stringifyRegex(input: RegExp | string): string {
   if (typeof input === 'string')
     return input
   return `$$vitest:${input.toString()}`
 }
 
-function getTestConfig(ctx: Vitest) {
+function getTestConfig(ctx: Vitest): ResolvedConfig {
   const config = ctx.getSerializableConfig()
   // v8 serialize does not support regex
-  return {
+  return <ResolvedConfig>{
     ...config,
     testNamePattern: config.testNamePattern
       ? stringifyRegex(config.testNamePattern)
@@ -83,7 +83,7 @@ export function createChildProcessPool(ctx: Vitest, { execArgv, env }: PoolProce
     })
   }
 
-  async function runWithFiles(files: string[], invalidates: string[] = []) {
+  async function runWithFiles(files: string[], invalidates: string[] = []): Promise<void> {
     ctx.state.clearFiles(files)
     const config = getTestConfig(ctx)
 
@@ -119,6 +119,7 @@ export function createChildProcessPool(ctx: Vitest, { execArgv, env }: PoolProce
         if (!child.killed)
           child.kill()
       })
+      children.clear()
     },
   }
 }

packages/vitest/src/types/config.ts

@@ -16,6 +16,7 @@ export type { SequenceHooks, SequenceSetupFiles } from '@vitest/runner'
 export type BuiltinEnvironment = 'node' | 'jsdom' | 'happy-dom' | 'edge-runtime'
 // Record is used, so user can get intellisense for builtin environments, but still allow custom environments
 export type VitestEnvironment = BuiltinEnvironment | (string & Record<never, never>)
+export type VitestPool = 'threads' | 'child_process'
 export type CSSModuleScopeStrategy = 'stable' | 'scoped' | 'non-scoped'
 
 export type ApiConfig = Pick<CommonServerOptions, 'port' | 'strictPort' | 'host'>
@@ -162,6 +163,21 @@ export interface InlineConfig {
    */
   environmentMatchGlobs?: [string, VitestEnvironment][]
 
+  /**
+   * Automatically assign pool based on globs. The first match will be used.
+   *
+   * Format: [glob, pool-name]
+   *
+   * @default []
+   * @example [
+   *   // all tests in "browser" directory will run in an actual browser
+   *   ['tests/browser/**', 'browser'],
+   *   // all other tests will run based on "threads" option, if you didn't specify other globs
+   *   // ...
+   * ]
+   */
+  poolMatchGlobs?: [string, VitestPool][]
+
   /**
    * Update snapshot
    *

packages/vitest/src/utils/test-helpers.ts

@@ -25,7 +25,7 @@ export async function groupFilesByEnv(files: string[], config: ResolvedConfig) {
     // 2. Check for globals
     if (!env) {
       for (const [glob, target] of config.environmentMatchGlobs || []) {
-        if (mm.isMatch(file, glob)) {
+        if (mm.isMatch(file, glob, { cwd: config.root })) {
           env = target
           break
         }

pnpm-lock.yaml

@@ -0,0 +1,12 @@
+{
+  "name": "@vitest/test-mixed-pools",
+  "type": "module",
+  "private": true,
+  "scripts": {
+    "test": "vitest",
+    "coverage": "vitest run --coverage"
+  },
+  "devDependencies": {
+    "vitest": "workspace:*"
+  }
+}

test/mixed-pools/package.json

@@ -0,0 +1,11 @@
+import { isMainThread, threadId } from 'node:worker_threads'
+import { expect, test } from 'vitest'
+
+test('has access to child_process API', () => {
+  expect(process.send).toBeDefined()
+})
+
+test('doesn\'t have access to threads API', () => {
+  expect(isMainThread).toBe(true)
+  expect(threadId).toBe(0)
+})

test/mixed-pools/test/child-specific.child_process.test.ts

@@ -0,0 +1,11 @@
+import { isMainThread, threadId } from 'node:worker_threads'
+import { expect, test } from 'vitest'
+
+test('has access access to worker API', () => {
+  expect(isMainThread).toBe(false)
+  expect(threadId).toBeGreaterThan(0)
+})
+
+test('doesn\'t have access access to child_process API', () => {
+  expect(process.send).toBeUndefined()
+})

test/mixed-pools/test/threads-specific.threads.test.ts

@@ -0,0 +1,11 @@
+import { defineConfig } from 'vite'
+
+export default defineConfig({
+  test: {
+    include: ['test/*.test.ts'],
+    poolMatchGlobs: [
+      ['**/test/*.child_process.test.ts', 'child_process'],
+      ['**/test/*.threads.test.ts', 'threads'],
+    ],
+  },
+})