fix(snapshot): improve inline snapshot inside loop rejection (#6339)

Author: hi-ogawa

packages/snapshot/src/port/state.ts

@@ -53,6 +53,7 @@ export default class SnapshotState {
   private _snapshotData: SnapshotData
   private _initialData: SnapshotData
   private _inlineSnapshots: Array<InlineSnapshot>
+  private _inlineSnapshotStacks: Array<ParsedStack>
   private _rawSnapshots: Array<RawSnapshot>
   private _uncheckedKeys: Set<string>
   private _snapshotFormat: PrettyFormatOptions
@@ -77,6 +78,7 @@ export default class SnapshotState {
     this._snapshotData = data
     this._dirty = dirty
     this._inlineSnapshots = []
+    this._inlineSnapshotStacks = []
     this._rawSnapshots = []
     this._uncheckedKeys = new Set(Object.keys(this._snapshotData))
     this._counters = new Map()
@@ -136,38 +138,13 @@ export default class SnapshotState {
   private _addSnapshot(
     key: string,
     receivedSerialized: string,
-    options: { isInline: boolean; rawSnapshot?: RawSnapshotInfo; error?: Error },
+    options: { rawSnapshot?: RawSnapshotInfo; stack?: ParsedStack },
   ): void {
     this._dirty = true
-    if (options.isInline) {
-      const error = options.error || new Error('snapshot')
-      const stacks = parseErrorStacktrace(
-        error,
-        { ignoreStackEntries: [] },
-      )
-      const _stack = this._inferInlineSnapshotStack(stacks)
-      if (!_stack) {
-        throw new Error(
-          `@vitest/snapshot: Couldn't infer stack frame for inline snapshot.\n${JSON.stringify(
-            stacks,
-          )}`,
-        )
-      }
-      const stack = this.environment.processStackTrace?.(_stack) || _stack
-      // removing 1 column, because source map points to the wrong
-      // location for js files, but `column-1` points to the same in both js/ts
-      // https://github.com/vitejs/vite/issues/8657
-      stack.column--
-      // reject multiple inline snapshots at the same location
-      const duplicateIndex = this._inlineSnapshots.findIndex(s => s.file === stack.file && s.line === stack.line && s.column === stack.column)
-      if (duplicateIndex >= 0) {
-        // remove the first one to avoid updating an inline snapshot
-        this._inlineSnapshots.splice(duplicateIndex, 1)
-        throw new Error('toMatchInlineSnapshot cannot be called multiple times at the same location.')
-      }
+    if (options.stack) {
       this._inlineSnapshots.push({
         snapshot: receivedSerialized,
-        ...stack,
+        ...options.stack,
       })
     }
     else if (options.rawSnapshot) {
@@ -316,6 +293,36 @@ export default class SnapshotState {
       this._snapshotData[key] = receivedSerialized
     }
 
+    // find call site of toMatchInlineSnapshot
+    let stack: ParsedStack | undefined
+    if (isInline) {
+      const stacks = parseErrorStacktrace(
+        error || new Error('snapshot'),
+        { ignoreStackEntries: [] },
+      )
+      const _stack = this._inferInlineSnapshotStack(stacks)
+      if (!_stack) {
+        throw new Error(
+          `@vitest/snapshot: Couldn't infer stack frame for inline snapshot.\n${JSON.stringify(
+            stacks,
+          )}`,
+        )
+      }
+      stack = this.environment.processStackTrace?.(_stack) || _stack
+      // removing 1 column, because source map points to the wrong
+      // location for js files, but `column-1` points to the same in both js/ts
+      // https://github.com/vitejs/vite/issues/8657
+      stack.column--
+
+      // reject multiple inline snapshots at the same location
+      if (this._inlineSnapshotStacks.some(s => s.file === stack!.file && s.line === stack!.line && s.column === stack!.column)) {
+        // remove already succeeded snapshot
+        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)
+    }
+
     // These are the conditions on when to write snapshots:
     //  * There's no snapshot file in a non-CI environment.
     //  * There is a snapshot file and we decided to update the snapshot.
@@ -338,8 +345,7 @@ export default class SnapshotState {
           }
 
           this._addSnapshot(key, receivedSerialized, {
-            error,
-            isInline,
+            stack,
             rawSnapshot,
           })
         }
@@ -349,8 +355,7 @@ export default class SnapshotState {
       }
       else {
         this._addSnapshot(key, receivedSerialized, {
-          error,
-          isInline,
+          stack,
           rawSnapshot,
         })
         this.added++

test/cli/fixtures/fails/inline-snapshop-inside-loop-update-none.test.ts

@@ -1,22 +1,30 @@
 import {test, expect} from "vitest";
 
-test("ok", () => {
-  expect("ok").toMatchInlineSnapshot(`"ok"`);
-});
-
 test("fail 1", () => {
   for (const str of ["foo", "bar"]) {
     expect(str).toMatchInlineSnapshot();
   }
 });
 
+test("fail 2.1", () => {
+  for (const str of ["foo", "bar"]) {
+    expect(str).toMatchInlineSnapshot(`"foo"`);
+  }
+});
+
+test("fail 2.2", () => {
+  for (const str of ["foo", "bar"]) {
+    expect(str).toMatchInlineSnapshot(`"bar"`);
+  }
+});
+
 test("fail 3", () => {
   for (const str of ["ok", "ok"]) {
     expect(str).toMatchInlineSnapshot();
   }
 });
 
-test("somehow ok", () => {
+test("fail 4", () => {
   for (const str of ["ok", "ok"]) {
     expect(str).toMatchInlineSnapshot(`"ok"`);
   }

test/cli/fixtures/fails/inline-snapshop-inside-loop-update-all.test.ts renamed to test/cli/fixtures/fails/inline-snapshop-inside-loop.test.ts

@@ -31,13 +31,14 @@ Error: InlineSnapshot cannot be used inside of test.each or describe.each
 Error: InlineSnapshot cannot be used inside of test.each or describe.each"
 `;
 
-exports[`should fail inline-snapshop-inside-loop-update-all.test.ts > inline-snapshop-inside-loop-update-all.test.ts 1`] = `
+exports[`should fail inline-snapshop-inside-loop.test.ts > inline-snapshop-inside-loop.test.ts 1`] = `
 "Error: toMatchInlineSnapshot cannot be called multiple times at the same location.
+Error: toMatchInlineSnapshot cannot be called multiple times at the same location.
+Error: toMatchInlineSnapshot cannot be called multiple times at the same location.
+Error: toMatchInlineSnapshot cannot be called multiple times at the same location.
 Error: toMatchInlineSnapshot cannot be called multiple times at the same location."
 `;
 
-exports[`should fail inline-snapshop-inside-loop-update-none.test.ts > inline-snapshop-inside-loop-update-none.test.ts 1`] = `"Error: Snapshot \`fail 2 1\` mismatched"`;
-
 exports[`should fail mock-import-proxy-module.test.ts > mock-import-proxy-module.test.ts 1`] = `"Error: There are some problems in resolving the mocks API."`;
 
 exports[`should fail nested-suite.test.ts > nested-suite.test.ts 1`] = `"AssertionError: expected true to be false // Object.is equality"`;

test/cli/test/__snapshots__/fails.test.ts.snap

@@ -10,7 +10,7 @@ const files = await fg('**/*.test.ts', { cwd: root, dot: true })
 it.each(files)('should fail %s', async (file) => {
   const { stderr } = await runVitest({
     root,
-    update: file === 'inline-snapshop-inside-loop-update-all.test.ts' ? true : undefined,
+    update: file === 'inline-snapshop-inside-loop.test.ts' ? true : undefined,
   }, [file])
 
   expect(stderr).toBeTruthy()

test/cli/test/fails.test.ts

@@ -1,4 +1,3 @@
-/* eslint-disable ts/prefer-ts-expect-error */
 /* eslint-disable ts/ban-ts-comment */
 
 import { google, type sheets_v4 } from 'googleapis'