refactor: proposal for the public Vitest API

Author: sheremet-va

packages/browser/src/node/pool.ts

@@ -1,7 +1,7 @@
 import * as nodeos from 'node:os'
 import crypto from 'node:crypto'
 import { relative } from 'pathe'
-import type { BrowserProvider, ProcessPool, Vitest, WorkspaceProject, WorkspaceSpec } from 'vitest/node'
+import type { BrowserProvider, ProcessPool, TestProject, TestSpecification, Vitest } from 'vitest/node'
 import { createDebugger } from 'vitest/node'
 
 const debug = createDebugger('vitest:browser:pool')
@@ -12,15 +12,15 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
   const waitForTests = async (
     method: 'run' | 'collect',
     contextId: string,
-    project: WorkspaceProject,
+    project: TestProject,
     files: string[],
   ) => {
     const context = project.browser!.state.createAsyncContext(method, contextId, files)
     return await context
   }
 
-  const executeTests = async (method: 'run' | 'collect', project: WorkspaceProject, files: string[]) => {
-    ctx.state.clearFiles(project, files)
+  const executeTests = async (method: 'run' | 'collect', project: TestProject, files: string[]) => {
+    ctx.state.clearFiles(project.workspaceProject, files)
     const browser = project.browser!
 
     const threadsCount = getThreadsCount(project)
@@ -33,7 +33,7 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
 
     if (!origin) {
       throw new Error(
-        `Can't find browser origin URL for project "${project.getName()}" when running tests for files "${files.join('", "')}"`,
+        `Can't find browser origin URL for project "${project.name}" when running tests for files "${files.join('", "')}"`,
       )
     }
 
@@ -50,7 +50,7 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
 
     debug?.(
       `[%s] Running %s tests in %s chunks (%s threads)`,
-      project.getName() || 'core',
+      project.name || 'core',
       files.length,
       chunks.length,
       threadsCount,
@@ -92,12 +92,12 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
     await Promise.all(promises)
   }
 
-  const runWorkspaceTests = async (method: 'run' | 'collect', specs: WorkspaceSpec[]) => {
-    const groupedFiles = new Map<WorkspaceProject, string[]>()
-    for (const [project, file] of specs) {
-      const files = groupedFiles.get(project) || []
-      files.push(file)
-      groupedFiles.set(project, files)
+  const runWorkspaceTests = async (method: 'run' | 'collect', specs: TestSpecification[]) => {
+    const groupedFiles = new Map<TestProject, string[]>()
+    for (const spec of specs) {
+      const files = groupedFiles.get(spec.project) || []
+      files.push(spec.moduleId)
+      groupedFiles.set(spec.project, files)
     }
 
     let isCancelled = false
@@ -120,7 +120,7 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
       ? nodeos.availableParallelism()
       : nodeos.cpus().length
 
-  function getThreadsCount(project: WorkspaceProject) {
+  function getThreadsCount(project: TestProject) {
     const config = project.config.browser
     if (!config.headless || !project.browser!.provider.supportsParallelism) {
       return 1

packages/vitest/src/node/context.ts

@@ -0,0 +1,25 @@
+import type { ProvidedContext } from '../types/general'
+import type { WorkspaceProject } from './workspace'
+
+export class VitestContext {
+  constructor(private workspaceProject: WorkspaceProject) {
+    this.workspaceProject = workspaceProject
+  }
+
+  /**
+   * Provide a custom serializable context to the project. This context will be available for tests once they run.
+   */
+  public provide<T extends keyof ProvidedContext & string>(
+    key: T,
+    value: ProvidedContext[T],
+  ): void {
+    this.workspaceProject.provide(key, value)
+  }
+
+  /**
+   * Get a custom serializable context provided to the project.
+   */
+  public get<T extends keyof ProvidedContext & string>(name: T): ProvidedContext[T] {
+    return this.workspaceProject.getProvidedContext()[name]
+  }
+}

packages/vitest/src/node/core.ts

@@ -668,7 +668,7 @@ export class Vitest {
         process.exitCode = 1
       }
     })()
-      .finally(async () => {
+      .finally(() => {
         this.runningPromise = undefined
 
         // all subsequent runs will treat this as a fresh run
@@ -1054,7 +1054,9 @@ export class Vitest {
           else if (runningServers > 1) {
             console.warn(`Tests closed successfully but something prevents ${runningServers} Vite servers from exiting`)
           }
-          else { console.warn('Tests closed successfully but something prevents the main process from exiting') }
+          else {
+            console.warn('Tests closed successfully but something prevents the main process from exiting')
+          }
 
           console.warn('You can try to identify the cause by enabling "hanging-process" reporter. See https://vitest.dev/config/#reporters')
         }

packages/vitest/src/node/publicVitest.ts

@@ -0,0 +1,48 @@
+/* eslint-disable ts/method-signature-style */
+import type { ModuleNode, ViteDevServer } from 'vite'
+import { SnapshotManager } from '@vitest/snapshot/manager'
+import type { TestProject } from './reported-workspace-project'
+import type { TestModule } from './reporters/reported-tasks'
+import type { TestSpecification } from './spec'
+import type { ResolvedConfig } from './types/config'
+import type { VitestContext } from './context'
+
+export interface Vitest {
+  readonly config: ResolvedConfig
+  readonly projects: TestProject[]
+  readonly vite: ViteDevServer
+  readonly moduleGraph: VitestModuleGraph
+  readonly runner: VitestRunner
+  readonly context: VitestContext
+  readonly snapshot: VitestSnapshot
+
+  onClose(cb: () => void): void
+
+  close(): Promise<void>
+  exit(): Promise<void>
+}
+
+interface VitestModuleGraph {
+  getTestModuleIds(moduleNames?: string[]): string[]
+  getViteModuleNodeById(
+    transformMode: 'web' | 'ssr' | 'browser',
+    moduleId: string,
+  ): ModuleNode | undefined
+  getViteModuleNodesById(moduleId: string): ModuleNode[]
+}
+
+interface VitestRunner {
+  // Vitest starts a standalone runner, will react on watch changes, it doesn't run tests
+  start(): Promise<void>
+
+  // Vitest will still start in a standalone mode if `watch` is `true`
+  run(): Promise<TestModule[]>
+  runModules(moduleNames: string[]): Promise<TestModule[]>
+  runTests(filters: TestSpecification[]): Promise<TestModule[]>
+}
+
+class VitestSnapshot extends SnapshotManager {
+  async update(_files?: string[]): Promise<void> {
+    // TODO
+  }
+}

packages/vitest/src/node/reported-workspace-project.ts

@@ -2,6 +2,9 @@ import type { ProvidedContext } from '../types/general'
 import type { ResolvedConfig, ResolvedProjectConfig, SerializedConfig } from './types/config'
 import type { WorkspaceProject } from './workspace'
 import type { Vitest } from './core'
+import type { BrowserServer } from './types/browser'
+import { type WorkspaceSpec, getFilePoolName } from './pool'
+import { VitestContext } from './context'
 
 export class TestProject {
   /**
@@ -29,43 +32,51 @@ export class TestProject {
    */
   public readonly name: string
 
+  /**
+   * Context manager.
+   */
+  public readonly context: VitestContext
+
   constructor(workspaceProject: WorkspaceProject) {
     this.workspaceProject = workspaceProject
     this.vitest = workspaceProject.ctx
     this.globalConfig = workspaceProject.ctx.config
     this.config = workspaceProject.config
     this.name = workspaceProject.getName()
+    this.context = new VitestContext(workspaceProject)
   }
 
   /**
    * Serialized project configuration. This is the config that tests receive.
    */
-  public get serializedConfig() {
+  public get serializedConfig(): SerializedConfig {
     return this.workspaceProject.getSerializableConfig()
   }
 
   /**
-   * Custom context provided to the project.
+   * Browser server if the project has browser enabled.
    */
-  public context(): ProvidedContext {
-    return this.workspaceProject.getProvidedContext()
+  public get browser(): BrowserServer | null {
+    return this.workspaceProject.browser || null
   }
 
   /**
-   * Provide a custom serializable context to the project. This context will be available for tests once they run.
+   * Create test specification describing a test module.
+   * @param moduleId File path to the module or an ID that Vite understands (like a virtual module).
+   * @param pool The pool to run the test in. If not provided, a pool will be selected based on the project configuration.
    */
-  public provide<T extends keyof ProvidedContext & string>(
-    key: T,
-    value: ProvidedContext[T],
-  ): void {
-    this.workspaceProject.provide(key, value)
+  public createSpecification(moduleId: string, pool?: string): WorkspaceSpec {
+    return this.workspaceProject.createSpec(
+      moduleId,
+      pool || getFilePoolName(this.workspaceProject, moduleId),
+    )
   }
 
   public toJSON(): SerializedTestProject {
     return {
       name: this.name,
       serializedConfig: this.serializedConfig,
-      context: this.context(),
+      context: this.workspaceProject.getProvidedContext(),
     }
   }
 }