feat: add --shard command (#1477)

* feat: add --shard command

* chore: increase timeout for shard tests

* chore: use relative path for spec

* chore: add try/catch to plugin

docs/guide/cli.md

@@ -66,14 +66,37 @@ vitest related /src/index.ts /src/hello-world.js
 | `--environment <env>` | Runner environment (default: `node`) |
 | `--passWithNoTests` | Pass when no tests found |
 | `--allowOnly` | Allow tests and suites that are marked as `only` (default: false in CI, true otherwise) |
-| `--changed [since]` | Run tests that are affected by the changed files (default: false). See [docs](#changed)
+| `--changed [since]` | Run tests that are affected by the changed files (default: false). See [docs](#changed) |
+| `--shard <shard>` | Execute tests in a specified shard |
 | `-h, --help` | Display available CLI options |
 
 ### changed
 
 - **Type**: `boolean | string`
 - **Default**: false
 
-Run tests only against changed files. If no value is provided, it will run tests against uncommitted changes (including staged and unstaged).
+  Run tests only against changed files. If no value is provided, it will run tests against uncommitted changes (including staged and unstaged).
 
-To run tests against changes made in the last commit, you can use `--changed HEAD~1`. You can also pass commit hash or branch name.
+  To run tests against changes made in the last commit, you can use `--changed HEAD~1`. You can also pass commit hash or branch name.
+
+### shard
+
+- **Type**: `string`
+- **Default**: disabled
+
+  Test suite shard to execute in a format of `<index>`/`<count>`, where
+
+  - `count` is a positive integer, count of divided parts
+  - `index` is a positive integer, index of divided part
+
+  This command will divide all tests into `count` equal parts, and will run only those that happen to be in an `index` part. For example, to split your tests suite into three parts, use this:
+
+  ```sh
+  vitest run --shard=1/3
+  vitest run --shard=2/3
+  vitest run --shard=3/3
+  ```
+
+:::warning
+You cannot use this option with `--watch` enabled (enabled in dev by default).
+:::

packages/vitest/src/node/cli.ts

@@ -33,6 +33,7 @@ cli
   .option('--environment <env>', 'runner environment (default: node)')
   .option('--passWithNoTests', 'pass when no tests found')
   .option('--allowOnly', 'Allow tests and suites that are marked as only (default: !process.env.CI)')
+  .option('--shard <shard>', 'Test suite shard to execute in a format of <index>/<count>')
   .option('--changed [since]', 'Run tests that are affected by the changed files (default: false)')
   .help()
 

packages/vitest/src/node/config.ts

@@ -90,6 +90,23 @@ export function resolveConfig(
 
   resolved.coverage = resolveC8Options(options.coverage || {}, resolved.root)
 
+  if (options.shard) {
+    if (resolved.watch)
+      throw new Error('You cannot use --shard option with enabled watch')
+
+    const [indexString, countString] = options.shard.split('/')
+    const index = Math.abs(parseInt(indexString, 10))
+    const count = Math.abs(parseInt(countString, 10))
+
+    if (isNaN(count) || count <= 0)
+      throw new Error('--shard <count> must be a positive number')
+
+    if (isNaN(index) || index <= 0 || index > count)
+      throw new Error('--shard <index> must be a positive number less then <count>')
+
+    resolved.shard = { index, count }
+  }
+
   resolved.deps = resolved.deps || {}
   // vitenode will try to import such file with native node,
   // but then our mocker will not work properly

packages/vitest/src/node/plugins/index.ts

@@ -125,10 +125,16 @@ export async function VitestPlugin(options: UserConfig = {}, ctx = new Vitest())
       async configureServer(server) {
         if (haveStarted)
           await ctx.report('onServerRestart')
-        await ctx.setServer(options, server)
-        haveStarted = true
-        if (options.api && options.watch)
-          (await import('../../api/setup')).setup(ctx)
+        try {
+          await ctx.setServer(options, server)
+          haveStarted = true
+          if (options.api && options.watch)
+            (await import('../../api/setup')).setup(ctx)
+        }
+        catch (err) {
+          ctx.printError(err, true)
+          process.exit(1)
+        }
 
         // #415, in run mode we don't need the watcher, close it would improve the performance
         if (!options.watch)

packages/vitest/src/node/pool.ts

@@ -1,14 +1,15 @@
 import { MessageChannel } from 'worker_threads'
 import { pathToFileURL } from 'url'
 import { cpus } from 'os'
+import { createHash } from 'crypto'
 import { resolve } from 'pathe'
 import type { Options as TinypoolOptions } from 'tinypool'
 import { Tinypool } from 'tinypool'
 import { createBirpc } from 'birpc'
 import type { RawSourceMap } from 'vite-node'
 import type { WorkerContext, WorkerRPC } from '../types'
 import { distDir } from '../constants'
-import { AggregateError } from '../utils'
+import { AggregateError, slash } from '../utils'
 import type { Vitest } from './core'
 
 export type RunWithFiles = (files: string[], invalidates?: string[]) => Promise<void>
@@ -88,6 +89,27 @@ export function createPool(ctx: Vitest): WorkerPool {
     }
 
     return async (files, invalidates) => {
+      if (config.shard) {
+        const { index, count } = config.shard
+        const shardSize = Math.ceil(files.length / count)
+        const shardStart = shardSize * (index - 1)
+        const shardEnd = shardSize * index
+        files = files
+          .map((file) => {
+            const fullPath = resolve(slash(config.root), slash(file))
+            const specPath = fullPath.slice(config.root.length)
+            return {
+              file,
+              hash: createHash('sha1')
+                .update(specPath)
+                .digest('hex'),
+            }
+          })
+          .sort((a, b) => (a.hash < b.hash ? -1 : a.hash > b.hash ? 1 : 0))
+          .slice(shardStart, shardEnd)
+          .map(({ file }) => file)
+      }
+
       if (!ctx.config.threads) {
         await runFiles(files)
       }

packages/vitest/src/types/config.ts

@@ -386,9 +386,17 @@ export interface UserConfig extends InlineConfig {
    * @default false
    */
   changed?: boolean | string
+
+  /**
+   * Test suite shard to execute in a format of <index>/<count>.
+   * Will divide tests into a `count` numbers, and run only the `indexed` part.
+   * Cannot be used with enabled watch.
+   * @example --shard=2/3
+   */
+  shard?: string
 }
 
-export interface ResolvedConfig extends Omit<Required<UserConfig>, 'config' | 'filters' | 'coverage' | 'testNamePattern' | 'related' | 'api' | 'reporters' | 'resolveSnapshotPath'> {
+export interface ResolvedConfig extends Omit<Required<UserConfig>, 'config' | 'filters' | 'coverage' | 'testNamePattern' | 'related' | 'api' | 'reporters' | 'resolveSnapshotPath' | 'shard'> {
   base?: string
 
   config?: string
@@ -404,4 +412,8 @@ export interface ResolvedConfig extends Omit<Required<UserConfig>, 'config' | 'f
   defines: Record<string, any>
 
   api?: ApiConfig
+  shard?: {
+    index: number
+    count: number
+  }
 }

test/shard/package.json

@@ -0,0 +1,11 @@
+{
+  "name": "@vitest/test-shard",
+  "private": true,
+  "scripts": {
+    "test": "vitest run shard-test.test.ts"
+  },
+  "devDependencies": {
+    "execa": "^6.1.0",
+    "vitest": "workspace:*"
+  }
+}

test/shard/shard-test.test.ts

@@ -0,0 +1,50 @@
+import { expect, test } from 'vitest'
+import { basename } from 'pathe'
+import { execa } from 'execa'
+
+const runVitest = async (args: string[]) => {
+  const { stdout } = await execa('vitest', ['--run', '--dir', './test', ...args])
+  return stdout
+}
+
+const parsePaths = (stdout: string) => {
+  return stdout
+    .split('\n')
+    .filter(line => line && line.includes('.test.js'))
+    .map(file => basename(file.trim().split(' ')[1]))
+    .sort()
+}
+
+test('--shard=1/1', async () => {
+  const stdout = await runVitest(['--shard=1/1'])
+
+  const paths = parsePaths(stdout)
+
+  expect(paths).toEqual(['1.test.js', '2.test.js', '3.test.js'])
+})
+
+test('--shard=1/2', async () => {
+  const stdout = await runVitest(['--shard=1/2'])
+
+  const paths = parsePaths(stdout)
+
+  expect(paths).toEqual(['1.test.js', '2.test.js'])
+})
+
+test('--shard=2/2', async () => {
+  const stdout = await runVitest(['--shard=2/2'])
+
+  const paths = parsePaths(stdout)
+
+  expect(paths).toEqual(['3.test.js'])
+})
+
+test('--shard=4/4', async () => {
+  const stdout = await runVitest(['--shard=4/4'])
+
+  const paths = parsePaths(stdout)
+
+  // project only has 3 files
+  // shards > 3 are empty
+  expect(paths).toEqual([])
+})

test/shard/test/1.test.js

@@ -0,0 +1,3 @@
+import { expect, test } from 'vitest'
+
+test('1', () => expect(1).toBe(1))

test/shard/test/2.test.js

@@ -0,0 +1,3 @@
+import { expect, test } from 'vitest'
+
+test('2', () => expect(1).toBe(1))

test/shard/test/3.test.js

@@ -0,0 +1,3 @@
+import { expect, test } from 'vitest'
+
+test('3', () => expect(1).toBe(1))

test/shard/vitest.config.ts

@@ -0,0 +1,7 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    testTimeout: 50_000,
+  },
+})