wip: state per test case

vitest-dev · Oct 31, 2024 · 20965e2 · 20965e2
1 parent 7c95367
commit 20965e2
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 50 deletions.
diff --git a/packages/snapshot/src/client.ts b/packages/snapshot/src/client.ts
@@ -51,67 +51,79 @@ export interface SnapshotClientOptions {
 }
 
 export class SnapshotClient {
-  snapshotStateMap: Map<string, SnapshotState> = new Map()
+  // snapshotStateMap: Map<string, SnapshotState> = new Map()
 
   constructor(private options: SnapshotClientOptions = {}) {}
 
-  async setup(
-    filepath: string,
-    options: SnapshotStateOptions,
-  ): Promise<void> {
-    if (this.snapshotStateMap.has(filepath)) {
-      throw new Error('already setup')
-    }
-    this.snapshotStateMap.set(
-      filepath,
-      await SnapshotState.create(filepath, options),
-    )
-  }
+  // async setup(
+  //   filepath: string,
+  //   options: SnapshotStateOptions,
+  // ): Promise<void> {
+  //   if (this.snapshotStateMap.has(filepath)) {
+  //     throw new Error('already setup')
+  //   }
+  //   this.snapshotStateMap.set(
+  //     filepath,
+  //     await SnapshotState.create(filepath, options),
+  //   )
+  // }
 
   async finish(filepath: string): Promise<SnapshotResult> {
-    const state = this.getSnapshotState(filepath)
-    const result = await state.pack()
-    this.snapshotStateMap.delete(filepath)
-    return result
+    const results: SnapshotResult[] = []
+    for (const testId of this.fileToTestIds.get(filepath)) {
+      const state = this.getSnapshotState(testId)
+      const result = await state.pack()
+      this.testIdToSnapshotPath.delete(testId)
+      results.push(result)
+    }
+    this.fileToTestIds.get(filepath)
+    // TODO: aggregate result
+    return results[0]
   }
 
   private fileToTestIds = new DefaultMap<string, Set<string>>(() => new Set())
-  private testIdToSnapshotPath: Record<string, string> = {}
+  private testIdToSnapshotPath = new Map<string, string>()
   private snapshotPathToState = new Map<string, SnapshotState>()
 
+  // TODO: concurrent safe
   async setupTest(
     filepath: string,
     testId: string,
     options: SnapshotStateOptions,
-  ): Promise<void> {
+  ): Promise<SnapshotState> {
     this.fileToTestIds.get(filepath).add(testId)
     const snapshotPath = await options.snapshotEnvironment.resolvePath(filepath)
-    this.testIdToSnapshotPath[testId] = snapshotPath
-    if (!this.snapshotPathToState.has(snapshotPath)) {
+    this.testIdToSnapshotPath.set(testId, snapshotPath)
+    // share same snapshot state for same snapshot path
+    let state = this.snapshotPathToState.get(snapshotPath)
+    if (!state) {
       const content = await options.snapshotEnvironment.readSnapshotFile(snapshotPath)
-      this.snapshotPathToState.set(
-        snapshotPath,
-        new SnapshotState(filepath, snapshotPath, content, options),
-      )
+      state = new SnapshotState(filepath, snapshotPath, content, options)
+      this.snapshotPathToState.set(snapshotPath, state)
     }
+    state.clearTest(testId)
+    return state
   }
 
-  skipTest(filepath: string, testName: string): void {
-    const state = this.getSnapshotState(filepath)
+  skipTest(testId: string, testName: string): void {
+    const state = this.getSnapshotState(testId)
     state.markSnapshotsAsCheckedForTest(testName)
   }
 
-  clearTest(filepath: string, testId: string): void {
-    const state = this.getSnapshotState(filepath)
-    state.clearTest(testId)
-  }
-
-  getSnapshotState(filepath: string): SnapshotState {
-    const state = this.snapshotStateMap.get(filepath)
-    if (!state) {
-      throw new Error('snapshot state not initialized')
+  // clearTest(testId: string): void {
+  //   const state = this.getSnapshotState(testId)
+  //   state.clearTest(testId)
+  // }
+
+  getSnapshotState(testId: string): SnapshotState {
+    const snapshotPath = this.testIdToSnapshotPath.get(testId)
+    if (snapshotPath) {
+      const state = this.snapshotPathToState.get(snapshotPath)
+      if (state) {
+        return state
+      }
     }
-    return state
+    throw new Error('snapshot state not initialized')
   }
 
   assert(options: AssertOptions): void {
@@ -133,7 +145,7 @@ export class SnapshotClient {
       throw new Error('Snapshot cannot be used outside of test')
     }
 
-    const snapshotState = this.getSnapshotState(filepath)
+    const snapshotState = this.getSnapshotState(testId)
 
     if (typeof properties === 'object') {
       if (typeof received !== 'object' || !received) {
@@ -197,7 +209,7 @@ export class SnapshotClient {
         throw new Error('Snapshot cannot be used outside of test')
       }
 
-      const snapshotState = this.getSnapshotState(filepath)
+      const snapshotState = this.getSnapshotState(options.testId)
 
       // save the filepath, so it don't lose even if the await make it out-of-context
       options.filepath ||= filepath
@@ -215,6 +227,9 @@ export class SnapshotClient {
   }
 
   clear(): void {
-    this.snapshotStateMap.clear()
+    this.fileToTestIds.clear()
+    this.testIdToSnapshotPath.clear()
+    this.snapshotPathToState.clear()
+    // this.snapshotStateMap.clear()
   }
 }
diff --git a/packages/vitest/src/runtime/runners/test.ts b/packages/vitest/src/runtime/runners/test.ts
@@ -59,7 +59,7 @@ export class VitestTestRunner implements VitestRunner {
       for (const test of getTests(suite)) {
         if (test.mode === 'skip') {
           const name = getNames(test).slice(1).join(' > ')
-          this.snapshotClient.skipTest(suite.file.filepath, name)
+          this.snapshotClient.skipTest(test.id, name)
         }
       }
 
@@ -102,18 +102,22 @@ export class VitestTestRunner implements VitestRunner {
     }
 
     // initialize snapshot state before running file suite
-    if (suite.mode !== 'skip' && 'filepath' in suite) {
-      await this.snapshotClient.setup(
-        suite.file.filepath,
-        this.workerState.config.snapshotOptions,
-      )
-    }
+    // if (suite.mode !== 'skip' && 'filepath' in suite) {
+    //   await this.snapshotClient.setup(
+    //     suite.file.filepath,
+    //     this.workerState.config.snapshotOptions,
+    //   )
+    // }
 
     this.workerState.current = suite
   }
 
-  onBeforeTryTask(test: Task) {
-    this.snapshotClient.clearTest(test.file.filepath, test.id)
+  async onBeforeTryTask(test: Task) {
+    const snapshotState = await this.snapshotClient.setupTest(
+      test.file.filepath,
+      test.id,
+      this.workerState.config.snapshotOptions,
+    )
     setState(
       {
         assertionCalls: 0,
@@ -123,7 +127,7 @@ export class VitestTestRunner implements VitestRunner {
         expectedAssertionsNumberErrorGen: null,
         testPath: test.file.filepath,
         currentTestName: getTestName(test),
-        snapshotState: this.snapshotClient.getSnapshotState(test.file.filepath),
+        snapshotState,
       },
       (globalThis as any)[GLOBAL_EXPECT],
     )