fix: ensure test context

vitest-dev · Nov 4, 2024 · 08db7bb · 08db7bb
1 parent 63d9817
commit 08db7bb
Showing 1 changed file with 15 additions and 7 deletions.
diff --git 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 `)
+    }
+    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',