fest(expect): improve predicate

vitest-dev · Jul 10, 2024 · 6b95014 · 6b95014
1 parent 992504b
commit 6b95014
Show file tree

Hide file tree

Showing 4 changed files with 244 additions and 63 deletions.
diff --git a/packages/expect/src/jest-expect.ts b/packages/expect/src/jest-expect.ts
@@ -699,17 +699,45 @@ export const JestChaiExpect: ChaiPlugin = (chai, utils) => {
       )
     },
   )
+
+  /**
+   * Used for `toHaveBeenCalledBefore` and `toHaveBeenCalledAfter` to determine if the expected spy was called before the result spy.
+   */
+  function predicate(expectInvocationCallOrder: number[], resultInvocationCallOrder: number[], failIfNoFirstInvocation: number): boolean {
+    if (expectInvocationCallOrder.length === 0) {
+      return !failIfNoFirstInvocation
+    }
+
+    if (resultInvocationCallOrder.length === 0) {
+      return false
+    }
+
+    return expectInvocationCallOrder[0] < resultInvocationCallOrder[0]
+  }
+
   def(
     ['toHaveBeenCalledBefore'],
-    function (resultSpy: MockInstance) {
+    function (resultSpy: MockInstance, failIfNoFirstInvocation = true) {
       const expectSpy = getSpy(this)
 
-      const [firstExpectSpyCall] = expectSpy.mock.invocationCallOrder
+      if (!isMockFunction(resultSpy)) {
+        throw new TypeError(
+          `${utils.inspect(resultSpy)} is not a spy or a call to a spy`,
+        )
+      }
 
-      const [firstResultSpyCall] = resultSpy.mock.invocationCallOrder
+      if (!isMockFunction(expectSpy)) {
+        throw new TypeError(
+          `${utils.inspect(expectSpy)} is not a spy or a call to a spy`,
+        )
+      }
 
       this.assert(
-        firstExpectSpyCall < firstResultSpyCall,
+        predicate(
+          expectSpy.mock.invocationCallOrder,
+          resultSpy.mock.invocationCallOrder,
+          failIfNoFirstInvocation,
+        ),
         `expected "${expectSpy.getMockName()}" to have been called before "${resultSpy.getMockName()}"`,
         `expected "${expectSpy.getMockName()}" to not have been called before "${resultSpy.getMockName()}"`,
         resultSpy,
@@ -719,15 +747,27 @@ export const JestChaiExpect: ChaiPlugin = (chai, utils) => {
   )
   def(
     ['toHaveBeenCalledAfter'],
-    function (resultSpy: MockInstance) {
+    function (resultSpy: MockInstance, failIfNoFirstInvocation = true) {
       const expectSpy = getSpy(this)
 
-      const [firstExpectSpyCall] = expectSpy.mock.invocationCallOrder
+      if (!isMockFunction(resultSpy)) {
+        throw new TypeError(
+          `${utils.inspect(resultSpy)} is not a spy or a call to a spy`,
+        )
+      }
 
-      const [firstResultSpyCall] = resultSpy.mock.invocationCallOrder
+      if (!isMockFunction(expectSpy)) {
+        throw new TypeError(
+          `${utils.inspect(expectSpy)} is not a spy or a call to a spy`,
+        )
+      }
 
       this.assert(
-        firstExpectSpyCall > firstResultSpyCall,
+        predicate(
+          resultSpy.mock.invocationCallOrder,
+          expectSpy.mock.invocationCallOrder,
+          failIfNoFirstInvocation,
+        ),
         `expected "${expectSpy.getMockName()}" to have been called after "${resultSpy.getMockName()}"`,
         `expected "${expectSpy.getMockName()}" to not have been called after "${resultSpy.getMockName()}"`,
         resultSpy,

diff --git a/packages/expect/src/types.ts b/packages/expect/src/types.ts
@@ -198,8 +198,8 @@ export interface Assertion<T = any>
   ) => void
   toHaveBeenCalledOnce: () => void
 
-  toHaveBeenCalledBefore: (mock: MockInstance) => void
-  toHaveBeenCalledAfter: (mock: MockInstance) => void
+  toHaveBeenCalledBefore: (mock: MockInstance, failIfNoFirstInvocation?: boolean) => void
+  toHaveBeenCalledAfter: (mock: MockInstance, failIfNoFirstInvocation?: boolean) => void
 
   toSatisfy: <E>(matcher: (value: E) => boolean, message?: string) => void