From d9b5f23db3041e03b1235c8a81f55ceb46dce8b5 Mon Sep 17 00:00:00 2001
From: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
Date: Mon, 4 Nov 2024 16:53:45 +0900
Subject: [PATCH] fix: ensure test context

---
 .../vitest/src/integrations/snapshot/chai.ts  | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/packages/vitest/src/integrations/snapshot/chai.ts b/packages/vitest/src/integrations/snapshot/chai.ts
index 9cd96fc2ba1b0..ef03d1f50b815 100644
--- a/packages/vitest/src/integrations/snapshot/chai.ts
+++ b/packages/vitest/src/integrations/snapshot/chai.ts
@@ -53,6 +53,14 @@ function getTestNames(test: Test) {
 }
 
 export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
+  function getTest(assertionName: string, obj: object) {
+    const test = utils.flag(obj, 'vitest-test')
+    if (!test) {
+      throw new Error(`'${assertionName}' cannot be used outside of test`)
+    }
+    return test as Test
+  }
+
   for (const key of ['matchSnapshot', 'toMatchSnapshot']) {
     utils.addMethod(
       chai.Assertion.prototype,
@@ -67,7 +75,7 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
           throw new Error(`${key} cannot be used with "not"`)
         }
         const expected = utils.flag(this, 'object')
-        const test: Test = utils.flag(this, 'vitest-test')
+        const test = getTest(key, this)
         if (typeof properties === 'string' && typeof message === 'undefined') {
           message = properties
           properties = undefined
@@ -94,7 +102,7 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
         throw new Error('toMatchFileSnapshot cannot be used with "not"')
       }
       const expected = utils.flag(this, 'object')
-      const test: Test = utils.flag(this, 'vitest-test')
+      const test = getTest('toMatchFileSnapshot', this)
       const errorMessage = utils.flag(this, 'message')
 
       const promise = getSnapshotClient().assertRaw({
@@ -125,8 +133,8 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
       if (isNot) {
         throw new Error('toMatchInlineSnapshot cannot be used with "not"')
       }
-      const test: Test = utils.flag(this, 'vitest-test')
-      const isInsideEach = test && (test.each || test.suite?.each)
+      const test = getTest('toMatchInlineSnapshot', this)
+      const isInsideEach = test.each || test.suite?.each
       if (isInsideEach) {
         throw new Error(
           'InlineSnapshot cannot be used inside of test.each or describe.each',
@@ -167,7 +175,7 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
         )
       }
       const expected = utils.flag(this, 'object')
-      const test: Test = utils.flag(this, 'vitest-test')
+      const test = getTest('toThrowErrorMatchingSnapshot', this)
       const promise = utils.flag(this, 'promise') as string | undefined
       const errorMessage = utils.flag(this, 'message')
       getSnapshotClient().assert({
@@ -192,8 +200,8 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
           'toThrowErrorMatchingInlineSnapshot cannot be used with "not"',
         )
       }
-      const test: Test = utils.flag(this, 'vitest-test')
-      const isInsideEach = test && (test.each || test.suite?.each)
+      const test = getTest('toThrowErrorMatchingInlineSnapshot', this)
+      const isInsideEach = test.each || test.suite?.each
       if (isInsideEach) {
         throw new Error(
           'InlineSnapshot cannot be used inside of test.each or describe.each',