fix: distribute test files to shards more evenly (#8288)

Shinyaigeek · AriPerkkio · web-flow · commit 7b489959a0a7 · 2025-07-12T10:10:44.000+03:00
Co-authored-by: Ari Perkkiö &lt;ari.perkkio@gmail.com&gt;
diff --git a/packages/vitest/src/node/sequencers/BaseSequencer.ts b/packages/vitest/src/node/sequencers/BaseSequencer.ts
@@ -16,9 +16,7 @@ export class BaseSequencer implements TestSequencer {
   public async shard(files: TestSpecification[]): Promise<TestSpecification[]> {
     const { config } = this.ctx
     const { index, count } = config.shard!
-    const shardSize = Math.ceil(files.length / count)
-    const shardStart = shardSize * (index - 1)
-    const shardEnd = shardSize * index
+    const [shardStart, shardEnd] = this.calculateShardRange(files.length, index, count)
     return [...files]
       .map((spec) => {
         const fullPath = resolve(slash(config.root), slash(spec.moduleId))
@@ -68,4 +66,20 @@ export class BaseSequencer implements TestSequencer {
       return bState.duration - aState.duration
     })
   }
+
+  // Calculate distributed shard range [start, end] distributed equally
+  private calculateShardRange(filesCount: number, index: number, count: number): [number, number] {
+    const baseShardSize = Math.floor(filesCount / count)
+    const remainderTestFilesCount = filesCount % count
+    if (remainderTestFilesCount >= index) {
+      const shardSize = baseShardSize + 1
+      const shardStart = shardSize * (index - 1)
+      const shardEnd = shardSize * index
+      return [shardStart, shardEnd]
+    }
+
+    const shardStart = remainderTestFilesCount * (baseShardSize + 1) + (index - remainderTestFilesCount - 1) * baseShardSize
+    const shardEnd = shardStart + baseShardSize
+    return [shardStart, shardEnd]
+  }
 }
diff --git a/test/config/fixtures/shard-4-files/test/1.test.js b/test/config/fixtures/shard-4-files/test/1.test.js
@@ -0,0 +1,5 @@
+import { test, expect } from 'vitest'
+
+test('test file 1', () => {
+  expect(1).toBe(1)
+})
diff --git a/test/config/fixtures/shard-4-files/test/2.test.js b/test/config/fixtures/shard-4-files/test/2.test.js
@@ -0,0 +1,5 @@
+import { test, expect } from 'vitest'
+
+test('test file 2', () => {
+  expect(2).toBe(2)
+})
diff --git a/test/config/fixtures/shard-4-files/test/3.test.js b/test/config/fixtures/shard-4-files/test/3.test.js
@@ -0,0 +1,5 @@
+import { test, expect } from 'vitest'
+
+test('test file 3', () => {
+  expect(3).toBe(3)
+})
diff --git a/test/config/fixtures/shard-4-files/test/4.test.js b/test/config/fixtures/shard-4-files/test/4.test.js
@@ -0,0 +1,5 @@
+import { test, expect } from 'vitest'
+
+test('test file 4', () => {
+  expect(4).toBe(4)
+})
diff --git a/test/config/fixtures/shard-4-files/vitest.config.js b/test/config/fixtures/shard-4-files/vitest.config.js
@@ -0,0 +1,7 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    include: ['test/**/*.test.js'],
+  },
+})
diff --git a/test/config/test/shard.test.ts b/test/config/test/shard.test.ts
@@ -4,8 +4,8 @@ import { basename } from 'pathe'
 import { expect, test } from 'vitest'
 import * as testUtils from '../../test-utils'
 
-function runVitest(config: TestUserConfig) {
-  return testUtils.runVitest({ ...config, root: './fixtures/shard' })
+function runVitest(config: TestUserConfig, root = './fixtures/shard') {
+  return testUtils.runVitest({ ...config, root })
 }
 
 function parsePaths(stdout: string) {
@@ -40,6 +40,52 @@ test('--shard=2/2', async () => {
   expect(paths).toEqual(['3.test.js'])
 })
 
+test('--shard=1/3 should distribute files evenly', async () => {
+  const { stdout } = await runVitest({ shard: '1/3' })
+
+  const paths = parsePaths(stdout)
+
+  // With 3 files and 3 shards, should get 1 file per shard
+  expect(paths).toEqual(['1.test.js'])
+})
+
+test('--shard=2/3 should distribute files evenly', async () => {
+  const { stdout } = await runVitest({ shard: '2/3' })
+
+  const paths = parsePaths(stdout)
+
+  // With 3 files and 3 shards, should get 1 file per shard
+  expect(paths).toEqual(['2.test.js'])
+})
+
+test('--shard=3/3 should distribute files evenly', async () => {
+  const { stdout } = await runVitest({ shard: '3/3' })
+
+  const paths = parsePaths(stdout)
+
+  // With 3 files and 3 shards, should get 1 file per shard
+  expect(paths).toEqual(['3.test.js'])
+})
+
+test('4 files with 3 shards should distribute evenly', async () => {
+  const { stdout: stdout1 } = await runVitest({ shard: '1/3' }, './fixtures/shard-4-files')
+  const { stdout: stdout2 } = await runVitest({ shard: '2/3' }, './fixtures/shard-4-files')
+  const { stdout: stdout3 } = await runVitest({ shard: '3/3' }, './fixtures/shard-4-files')
+
+  const paths1 = parsePaths(stdout1)
+  const paths2 = parsePaths(stdout2)
+  const paths3 = parsePaths(stdout3)
+
+  // Should distribute files more evenly: [2,1,1] instead of [2,2,0]
+  expect(paths1.length).toBe(2)
+  expect(paths2.length).toBe(1)
+  expect(paths3.length).toBe(1)
+
+  // All files should be covered exactly once
+  const allFiles = [...paths1, ...paths2, ...paths3].sort()
+  expect(allFiles).toEqual(['1.test.js', '2.test.js', '3.test.js', '4.test.js'])
+})
+
 test('--shard=4/4', async () => {
   const { stdout } = await runVitest({ shard: '4/4' })
 
diff --git a/test/core/test/sequencers.test.ts b/test/core/test/sequencers.test.ts
@@ -5,10 +5,11 @@ import { BaseSequencer } from '../../../packages/vitest/src/node/sequencers/Base
 import { RandomSequencer } from '../../../packages/vitest/src/node/sequencers/RandomSequencer'
 import { TestSpecification } from '../../../packages/vitest/src/node/spec'
 
-function buildCtx() {
+function buildCtx(config?: Partial<Vitest['config']>) {
   return {
     config: {
       sequence: {},
+      ...config,
     },
     cache: {
       getFileTestResults: vi.fn(),
@@ -128,6 +129,28 @@ describe('base sequencer', () => {
     const sorted = await sequencer.sort(files)
     expect(sorted).toStrictEqual(workspaced(['c', 'b', 'a']))
   })
+
+  test.each([
+    { files: 4, count: 3, expected: [2, 1, 1] },
+    { files: 5, count: 4, expected: [2, 1, 1, 1] },
+    { files: 9, count: 4, expected: [3, 2, 2, 2] },
+  ])('shard x/$count distributes $files files as $expected', async ({ count, files, expected }) => {
+    const specs = Array.from({ length: files }, (_, id) => ({ moduleId: `file-${id}.test.ts` } as TestSpecification))
+    const slices = []
+
+    for (const index of Array.from({ length: count }).keys()) {
+      const ctx = buildCtx({ root: '/example/root', shard: { index: 1 + index, count } })
+      const sequencer = new BaseSequencer(ctx)
+      const shard = await sequencer.shard(specs)
+
+      slices.push(shard.length)
+    }
+
+    expect(slices).toEqual(expected)
+
+    const sum = slices.reduce((total, current) => total + current, 0)
+    expect(sum).toBe(files)
+  })
 })
 
 describe('random sequencer', () => {
