fix(snapshot)!: reset snapshot state for retry and repeats

packages/snapshot/src/client.ts

@@ -34,8 +34,9 @@ export interface Context {
 
 interface AssertOptions {
   received: unknown
-  filepath?: string
-  name?: string
+  filepath: string
+  name: string
+  testId: string
   message?: string
   isInline?: boolean
   properties?: object
@@ -50,51 +51,53 @@ export interface SnapshotClientOptions {
 }
 
 export class SnapshotClient {
-  filepath?: string
-  name?: string
-  snapshotState: SnapshotState | undefined
   snapshotStateMap: Map<string, SnapshotState> = new Map()
 
   constructor(private options: SnapshotClientOptions = {}) {}
 
-  async startCurrentRun(
+  async setup(
     filepath: string,
-    name: string,
     options: SnapshotStateOptions,
   ): Promise<void> {
-    this.filepath = filepath
-    this.name = name
-
-    if (this.snapshotState?.testFilePath !== filepath) {
-      await this.finishCurrentRun()
-
-      if (!this.getSnapshotState(filepath)) {
-        this.snapshotStateMap.set(
-          filepath,
-          await SnapshotState.create(filepath, options),
-        )
-      }
-      this.snapshotState = this.getSnapshotState(filepath)
+    if (this.snapshotStateMap.has(filepath)) {
+      throw new Error('already setup')
     }
+    this.snapshotStateMap.set(
+      filepath,
+      await SnapshotState.create(filepath, options),
+    )
   }
 
-  getSnapshotState(filepath: string): SnapshotState {
-    return this.snapshotStateMap.get(filepath)!
+  async finish(filepath: string): Promise<SnapshotResult> {
+    const state = this.getSnapshotState(filepath)
+    const result = await state.pack()
+    this.snapshotStateMap.delete(filepath)
+    return result
+  }
+
+  skipTest(filepath: string, testName: string): void {
+    const state = this.getSnapshotState(filepath)
+    state.markSnapshotsAsCheckedForTest(testName)
   }
 
-  clearTest(): void {
-    this.filepath = undefined
-    this.name = undefined
+  clearTest(filepath: string, testId: string): void {
+    const state = this.getSnapshotState(filepath)
+    state.clearTest(testId)
   }
 
-  skipTestSnapshots(name: string): void {
-    this.snapshotState?.markSnapshotsAsCheckedForTest(name)
+  getSnapshotState(filepath: string): SnapshotState {
+    const state = this.snapshotStateMap.get(filepath)
+    if (!state) {
+      throw new Error('snapshot state not initialized')
+    }
+    return state
   }
 
   assert(options: AssertOptions): void {
     const {
-      filepath = this.filepath,
-      name = this.name,
+      filepath,
+      name,
+      testId,
       message,
       isInline = false,
       properties,
@@ -109,6 +112,8 @@ export class SnapshotClient {
       throw new Error('Snapshot cannot be used outside of test')
     }
 
+    const snapshotState = this.getSnapshotState(filepath)
+
     if (typeof properties === 'object') {
       if (typeof received !== 'object' || !received) {
         throw new Error(
@@ -122,7 +127,7 @@ export class SnapshotClient {
         if (!pass) {
           throw createMismatchError(
             'Snapshot properties mismatched',
-            this.snapshotState?.expand,
+            snapshotState.expand,
             received,
             properties,
           )
@@ -139,9 +144,8 @@ export class SnapshotClient {
 
     const testName = [name, ...(message ? [message] : [])].join(' > ')
 
-    const snapshotState = this.getSnapshotState(filepath)
-
     const { actual, expected, key, pass } = snapshotState.match({
+      testId,
       testName,
       received,
       isInline,
@@ -153,7 +157,7 @@ export class SnapshotClient {
     if (!pass) {
       throw createMismatchError(
         `Snapshot \`${key || 'unknown'}\` mismatched`,
-        this.snapshotState?.expand,
+        snapshotState.expand,
         actual?.trim(),
         expected?.trim(),
       )
@@ -165,7 +169,7 @@ export class SnapshotClient {
       throw new Error('Raw snapshot is required')
     }
 
-    const { filepath = this.filepath, rawSnapshot } = options
+    const { filepath, rawSnapshot } = options
 
     if (rawSnapshot.content == null) {
       if (!filepath) {
@@ -189,16 +193,6 @@ export class SnapshotClient {
     return this.assert(options)
   }
 
-  async finishCurrentRun(): Promise<SnapshotResult | null> {
-    if (!this.snapshotState) {
-      return null
-    }
-    const result = await this.snapshotState.pack()
-
-    this.snapshotState = undefined
-    return result
-  }
-
   clear(): void {
     this.snapshotStateMap.clear()
   }

packages/snapshot/src/port/inlineSnapshot.ts

@@ -9,6 +9,7 @@ import {
 
 export interface InlineSnapshot {
   snapshot: string
+  testId: string
   file: string
   line: number
   column: number

packages/snapshot/src/port/state.ts

@@ -23,6 +23,8 @@ import { saveRawSnapshots } from './rawSnapshot'
 
 import {
   addExtraLineBreaks,
+  CounterMap,
+  DefaultMap,
   getSnapshotData,
   keyToTestName,
   normalizeNewlines,
@@ -47,24 +49,24 @@ interface SaveStatus {
 }
 
 export default class SnapshotState {
-  private _counters: Map<string, number>
+  private _counters = new CounterMap<string>()
   private _dirty: boolean
   private _updateSnapshot: SnapshotUpdateState
   private _snapshotData: SnapshotData
   private _initialData: SnapshotData
   private _inlineSnapshots: Array<InlineSnapshot>
-  private _inlineSnapshotStacks: Array<ParsedStack>
+  private _inlineSnapshotStacks: Array<ParsedStack & { testId: string }>
+  private _testIdToKeys = new DefaultMap<string, string[]>(() => [])
   private _rawSnapshots: Array<RawSnapshot>
   private _uncheckedKeys: Set<string>
   private _snapshotFormat: PrettyFormatOptions
   private _environment: SnapshotEnvironment
   private _fileExists: boolean
-
-  added: number
+  private added = new CounterMap<string>()
+  private matched = new CounterMap<string>()
+  private unmatched = new CounterMap<string>()
+  private updated = new CounterMap<string>()
   expand: boolean
-  matched: number
-  unmatched: number
-  updated: number
 
   private constructor(
     public testFilePath: string,
@@ -74,20 +76,15 @@ export default class SnapshotState {
   ) {
     const { data, dirty } = getSnapshotData(snapshotContent, options)
     this._fileExists = snapshotContent != null // TODO: update on watch?
-    this._initialData = data
-    this._snapshotData = data
+    this._initialData = { ...data }
+    this._snapshotData = { ...data }
     this._dirty = dirty
     this._inlineSnapshots = []
     this._inlineSnapshotStacks = []
     this._rawSnapshots = []
     this._uncheckedKeys = new Set(Object.keys(this._snapshotData))
-    this._counters = new Map()
     this.expand = options.expand || false
-    this.added = 0
-    this.matched = 0
-    this.unmatched = 0
     this._updateSnapshot = options.updateSnapshot
-    this.updated = 0
     this._snapshotFormat = {
       printBasicPrototype: false,
       escapeString: false,
@@ -118,6 +115,31 @@ export default class SnapshotState {
     })
   }
 
+  clearTest(testId: string): void {
+    // clear inline
+    this._inlineSnapshots = this._inlineSnapshots.filter(s => s.testId !== testId)
+    this._inlineSnapshotStacks = this._inlineSnapshotStacks.filter(s => s.testId !== testId)
+
+    // clear file
+    for (const key of this._testIdToKeys.get(testId)) {
+      const name = keyToTestName(key)
+      const count = this._counters.get(name)
+      if (count > 0) {
+        if (key in this._snapshotData || key in this._initialData) {
+          this._snapshotData[key] = this._initialData[key]
+        }
+        this._counters.set(name, count - 1)
+      }
+    }
+    this._testIdToKeys.delete(testId)
+
+    // clear stats
+    this.added.delete(testId)
+    this.updated.delete(testId)
+    this.matched.delete(testId)
+    this.unmatched.delete(testId)
+  }
+
   protected _inferInlineSnapshotStack(stacks: ParsedStack[]): ParsedStack | null {
     // if called inside resolves/rejects, stacktrace is different
     const promiseIndex = stacks.findIndex(i =>
@@ -138,12 +160,13 @@ export default class SnapshotState {
   private _addSnapshot(
     key: string,
     receivedSerialized: string,
-    options: { rawSnapshot?: RawSnapshotInfo; stack?: ParsedStack },
+    options: { rawSnapshot?: RawSnapshotInfo; stack?: ParsedStack; testId: string },
   ): void {
     this._dirty = true
     if (options.stack) {
       this._inlineSnapshots.push({
         snapshot: receivedSerialized,
+        testId: options.testId,
         ...options.stack,
       })
     }
@@ -158,17 +181,6 @@ export default class SnapshotState {
     }
   }
 
-  clear(): void {
-    this._snapshotData = this._initialData
-    // this._inlineSnapshots = []
-    this._counters = new Map()
-    this.added = 0
-    this.matched = 0
-    this.unmatched = 0
-    this.updated = 0
-    this._dirty = false
-  }
-
   async save(): Promise<SaveStatus> {
     const hasExternalSnapshots = Object.keys(this._snapshotData).length
     const hasInlineSnapshots = this._inlineSnapshots.length
@@ -228,6 +240,7 @@ export default class SnapshotState {
   }
 
   match({
+    testId,
     testName,
     received,
     key,
@@ -236,12 +249,14 @@ export default class SnapshotState {
     error,
     rawSnapshot,
   }: SnapshotMatchOptions): SnapshotReturnOptions {
-    this._counters.set(testName, (this._counters.get(testName) || 0) + 1)
-    const count = Number(this._counters.get(testName))
+    // this also increments counter for inline snapshots. maybe we shouldn't?
+    this._counters.increment(testName)
+    const count = this._counters.get(testName)
 
     if (!key) {
       key = testNameToKey(testName, count)
     }
+    this._testIdToKeys.get(testId).push(key)
 
     // Do not mark the snapshot as "checked" if the snapshot is inline and
     // there's an external snapshot. This way the external snapshot can be
@@ -320,7 +335,7 @@ export default class SnapshotState {
         this._inlineSnapshots = this._inlineSnapshots.filter(s => !(s.file === stack!.file && s.line === stack!.line && s.column === stack!.column))
         throw new Error('toMatchInlineSnapshot cannot be called multiple times at the same location.')
       }
-      this._inlineSnapshotStacks.push(stack)
+      this._inlineSnapshotStacks.push({ ...stack, testId })
     }
 
     // These are the conditions on when to write snapshots:
@@ -338,27 +353,29 @@ export default class SnapshotState {
       if (this._updateSnapshot === 'all') {
         if (!pass) {
           if (hasSnapshot) {
-            this.updated++
+            this.updated.increment(testId)
           }
           else {
-            this.added++
+            this.added.increment(testId)
           }
 
           this._addSnapshot(key, receivedSerialized, {
             stack,
+            testId,
             rawSnapshot,
           })
         }
         else {
-          this.matched++
+          this.matched.increment(testId)
         }
       }
       else {
         this._addSnapshot(key, receivedSerialized, {
           stack,
+          testId,
           rawSnapshot,
         })
-        this.added++
+        this.added.increment(testId)
       }
 
       return {
@@ -371,7 +388,7 @@ export default class SnapshotState {
     }
     else {
       if (!pass) {
-        this.unmatched++
+        this.unmatched.increment(testId)
         return {
           actual: removeExtraLineBreaks(receivedSerialized),
           count,
@@ -384,7 +401,7 @@ export default class SnapshotState {
         }
       }
       else {
-        this.matched++
+        this.matched.increment(testId)
         return {
           actual: '',
           count,
@@ -415,10 +432,10 @@ export default class SnapshotState {
 
     const status = await this.save()
     snapshot.fileDeleted = status.deleted
-    snapshot.added = this.added
-    snapshot.matched = this.matched
-    snapshot.unmatched = this.unmatched
-    snapshot.updated = this.updated
+    snapshot.added = this.added.total()
+    snapshot.matched = this.matched.total()
+    snapshot.unmatched = this.unmatched.total()
+    snapshot.updated = this.updated.total()
     snapshot.unchecked = !status.deleted ? uncheckedCount : 0
     snapshot.uncheckedKeys = Array.from(uncheckedKeys)