refactor: add sequelizer

packages/vitest/src/node/cache/results.ts

@@ -3,7 +3,7 @@ import { dirname, resolve } from 'pathe'
 import type { File, ResolvedConfig } from '../../types'
 import { version } from '../../../package.json'
 
-interface SuiteResultCache {
+export interface SuiteResultCache {
   failed: boolean
   duration: number
 }

packages/vitest/src/node/pool.ts

@@ -1,16 +1,16 @@
 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 { ResolvedConfig, WorkerContext, WorkerRPC } from '../types'
 import { distDir } from '../constants'
-import { AggregateError, slash } from '../utils'
+import { AggregateError } from '../utils'
 import type { Vitest } from './core'
+import { BaseSequelizer } from './sequelizers/BaseSequelizer'
 
 export type RunWithFiles = (files: string[], invalidates?: string[]) => Promise<void>
 
@@ -86,54 +86,15 @@ export function createPool(ctx: Vitest): WorkerPool {
       }
     }
 
+    const sequelizer = new BaseSequelizer(ctx)
+
     return async (files, invalidates) => {
       const config = ctx.getSerializableConfig()
 
-      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)
-      }
-
-      files = files.sort((a, b) => {
-        const aState = ctx.state.results.getResults(a)
-        const bState = ctx.state.results.getResults(b)
-
-        if (!aState || !bState) {
-          const statsA = ctx.state.stats.getStats(a)
-          const statsB = ctx.state.stats.getStats(b)
-
-          if (!statsA || !statsB)
-            return !aState && bState ? -1 : !bState && aState ? 1 : 0
-
-          // run larger files first
-          return statsB.size - statsA.size
-        }
-
-        // run failed first
-        if (aState.failed && !bState.failed)
-          return -1
-        if (!aState.failed && bState.failed)
-          return 1
+      if (config.shard)
+        files = await sequelizer.shard(files)
 
-        // run longer first
-        return bState.duration - aState.duration
-      })
+      files = await sequelizer.sort(files)
 
       if (!ctx.config.threads) {
         await runFiles(config, files)

packages/vitest/src/node/sequelizers/BaseSequelizer.ts

@@ -0,0 +1,66 @@
+import { createHash } from 'crypto'
+import { resolve } from 'pathe'
+import { slash } from 'vite-node/utils'
+import type { Vitest } from '../core'
+import type { TestSequelizer } from './types'
+
+export class BaseSequelizer implements TestSequelizer {
+  protected ctx: Vitest
+
+  constructor(ctx: Vitest) {
+    this.ctx = ctx
+  }
+
+  // async so it can be extended by other sequelizers
+  public async shard(files: string[]): Promise<string[]> {
+    const config = this.ctx.getSerializableConfig()
+    const { index, count } = config.shard!
+    const shardSize = Math.ceil(files.length / count)
+    const shardStart = shardSize * (index - 1)
+    const shardEnd = shardSize * index
+    return [...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)
+  }
+
+  // async so it can be extended by other sequelizers
+  public async sort(files: string[]): Promise<string[]> {
+    const ctx = this.ctx
+    return [...files].sort((a, b) => {
+      const aState = ctx.state.getFileTestResults(a)
+      const bState = ctx.state.getFileTestResults(b)
+
+      if (!aState || !bState) {
+        const statsA = ctx.state.getFileStats(a)
+        const statsB = ctx.state.getFileStats(b)
+
+        // run unknown forst
+        if (!statsA || !statsB)
+          return !statsA && statsB ? -1 : !statsB && statsA ? 1 : 0
+
+        // run larger files first
+        return statsB.size - statsA.size
+      }
+
+      // run failed first
+      if (aState.failed && !bState.failed)
+        return -1
+      if (!aState.failed && bState.failed)
+        return 1
+
+      // run longer first
+      return bState.duration - aState.duration
+    })
+  }
+}

packages/vitest/src/node/sequelizers/types.ts

@@ -0,0 +1,10 @@
+import type { Awaitable } from '../../types'
+
+export interface TestSequelizer {
+  /**
+  * Slicing tests into shards. Will be run before `sort`.
+  * Only run, if `shard` is defined.
+  */
+  shard(files: string[]): Awaitable<string[]>
+  sort(files: string[]): Awaitable<string[]>
+}

packages/vitest/src/node/state.ts

@@ -10,6 +10,14 @@ export class StateManager {
   results = new ResultsCache()
   stats = new FilesCache()
 
+  getFileTestResults(id: string) {
+    return this.results.getResults(id)
+  }
+
+  getFileStats(id: string) {
+    return this.stats.getStats(id)
+  }
+
   catchError(err: unknown, type: string) {
     (err as ErrorWithDiff).type = type
     this.errorsSet.add(err)

test/core/test/sequelizers.test.ts

@@ -0,0 +1,97 @@
+import type { Vitest } from 'vitest'
+import { describe, expect, test, vi } from 'vitest'
+import { BaseSequelizer } from '../../../packages/vitest/src/node/sequelizers/BaseSequelizer'
+
+const buildCtx = () => {
+  return {
+    state: {
+      getFileTestResults: vi.fn(),
+      getFileStats: vi.fn(),
+    },
+  } as unknown as Vitest
+}
+
+describe('test sequelizers', () => {
+  test('sorting when no info is available', async () => {
+    const sequelizer = new BaseSequelizer(buildCtx())
+    const files = ['a', 'b', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(files)
+  })
+
+  test('prioritaze unknown files', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileStats').mockImplementation((file) => {
+      if (file === 'b')
+        return { size: 2 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['a', 'c', 'b'])
+  })
+
+  test('sort by size, larger first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileStats').mockImplementation((file) => {
+      if (file === 'a')
+        return { size: 1 }
+      if (file === 'b')
+        return { size: 2 }
+      if (file === 'c')
+        return { size: 3 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['c', 'b', 'a'])
+  })
+
+  test('sort by results, failed first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileTestResults').mockImplementation((file) => {
+      if (file === 'a')
+        return { failed: false, duration: 1 }
+      if (file === 'b')
+        return { failed: true, duration: 1 }
+      if (file === 'c')
+        return { failed: true, duration: 1 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['b', 'c', 'a'])
+  })
+
+  test('sort by results, long first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileTestResults').mockImplementation((file) => {
+      if (file === 'a')
+        return { failed: true, duration: 1 }
+      if (file === 'b')
+        return { failed: true, duration: 2 }
+      if (file === 'c')
+        return { failed: true, duration: 3 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['c', 'b', 'a'])
+  })
+
+  test('sort by results, long and failed first', async () => {
+    const ctx = buildCtx()
+    vi.spyOn(ctx.state, 'getFileTestResults').mockImplementation((file) => {
+      if (file === 'a')
+        return { failed: false, duration: 1 }
+      if (file === 'b')
+        return { failed: false, duration: 6 }
+      if (file === 'c')
+        return { failed: true, duration: 3 }
+    })
+    const sequelizer = new BaseSequelizer(ctx)
+    const files = ['b', 'a', 'c']
+    const sorted = await sequelizer.sort(files)
+    expect(sorted).toStrictEqual(['c', 'b', 'a'])
+  })
+})