From 38d469c4af9ca534095537dcdf25855d74b150b3 Mon Sep 17 00:00:00 2001
From: Benjamin Rupp <github@benjaminrupp.de>
Date: Tue, 6 Feb 2024 20:03:46 +0100
Subject: [PATCH] Add function that verifies that all when mocks are called

Inspired by verifyAllWhenMocksCalled of jest-when.
---
 README.md                | 19 +++++++++++++++++++
 src/behaviors.ts         |  4 ++++
 src/vitest-when.ts       | 14 +++++++++++++-
 test/vitest-when.test.ts | 21 ++++++++++++++++++++-
 4 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index cb40041..6048439 100644
--- a/README.md
+++ b/README.md
@@ -67,6 +67,25 @@ test('stubbing with vitest-when', () => {
 
 You should call `vi.resetAllMocks()` in your suite's `afterEach` hook to remove the implementation added by `when`. You can also set Vitest's [`mockReset`](https://vitest.dev/config/#mockreset) config to `true` instead of using `afterEach`.
 
+#### Supports verifying that all mocked functions were called
+
+Call `verifyAllWhenMocksCalled` after your test to assert that all mocks using `{ times: x }` were used.
+
+```javascript
+import { verifyAllWhenMocksCalled, when } from 'vitest-when'
+
+afterEach(() => {
+    verifyAllWhenMocksCalled() // fails because spy is only called once
+})
+
+it('should call spy two times', () => {
+    const spy = vi.fn()
+    when(spy, {times: 2}).calledWith(1, 2, 3).thenReturn(4)
+    spy(1, 2, 3)
+})
+```
+
+
 [vitest's mock functions]: https://vitest.dev/api/mock.html
 [stubs]: https://en.wikipedia.org/wiki/Test_stub
 [when]: #whenspy-tfunc-stubwrappertfunc
diff --git a/src/behaviors.ts b/src/behaviors.ts
index cc15c8a..d3adc72 100644
--- a/src/behaviors.ts
+++ b/src/behaviors.ts
@@ -8,6 +8,8 @@ export interface WhenOptions {
 export interface BehaviorStack<TFunc extends AnyFunction> {
   use: (args: Parameters<TFunc>) => BehaviorEntry<Parameters<TFunc>> | undefined
 
+  get: () => BehaviorEntry<Parameters<TFunc>>[]
+
   bindArgs: <TArgs extends Parameters<TFunc>>(
     args: TArgs,
     options: WhenOptions,
@@ -54,6 +56,8 @@ export const createBehaviorStack = <
       return behavior
     },
 
+    get: () => behaviors,
+
     bindArgs: (args, options) => ({
       addReturn: (values) => {
         behaviors.unshift(
diff --git a/src/vitest-when.ts b/src/vitest-when.ts
index a54a7db..dec04d9 100644
--- a/src/vitest-when.ts
+++ b/src/vitest-when.ts
@@ -1,10 +1,13 @@
+import { assert } from 'vitest'
 import { configureStub } from './stubs.ts'
-import type { WhenOptions } from './behaviors.ts'
+import type { BehaviorStack, WhenOptions } from './behaviors.ts'
 import type { AnyFunction } from './types.ts'
 
 export type { WhenOptions } from './behaviors.ts'
 export * from './errors.ts'
 
+export const behaviorStackRegistry = new Set<BehaviorStack<AnyFunction>>()
+
 export interface StubWrapper<TFunc extends AnyFunction> {
   calledWith<TArgs extends Parameters<TFunc>>(
     ...args: TArgs
@@ -25,6 +28,8 @@ export const when = <TFunc extends AnyFunction>(
 ): StubWrapper<TFunc> => {
   const behaviorStack = configureStub(spy)
 
+  behaviorStackRegistry.add(behaviorStack)
+
   return {
     calledWith: (...args) => {
       const behaviors = behaviorStack.bindArgs(args, options)
@@ -39,3 +44,10 @@ export const when = <TFunc extends AnyFunction>(
     },
   }
 }
+
+export const verifyAllWhenMocksCalled = () => {
+  const uncalledMocks = [...behaviorStackRegistry].flatMap((behaviorStack) => {
+    return behaviorStack.get().filter(behavior => behavior.times && behavior.times > 0)
+  })
+  assert.equal(uncalledMocks.length,0, `Failed verifyAllWhenMocksCalled: ${uncalledMocks.length} mock(s) not called:`)
+}
diff --git a/test/vitest-when.test.ts b/test/vitest-when.test.ts
index 4f3ae86..b94a44d 100644
--- a/test/vitest-when.test.ts
+++ b/test/vitest-when.test.ts
@@ -1,6 +1,7 @@
-import { vi, describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
 
 import * as subject from '../src/vitest-when.ts'
+import AssertionError = Chai.AssertionError
 
 declare module 'vitest' {
   interface AsymmetricMatchersContaining {
@@ -264,4 +265,22 @@ describe('vitest-when', () => {
     // intentionally do not call the spy
     expect(true).toBe(true)
   })
+
+  it('should fail if there are uncalled mocks and verify function is called', () => {
+    const spy = vi.fn()
+    subject.when(spy, { times: 2 }).calledWith(1, 2, 3).thenReturn(4)
+    spy(1, 2, 3)
+
+    let error: AssertionError | undefined = undefined
+    try {
+      subject.verifyAllWhenMocksCalled()
+    } catch (assertionError) {
+      error = assertionError as AssertionError
+      expect(error).toEqual(expect.objectContaining({ expected: 0, actual: 1 }))
+      expect(error.message).toContain(
+        'Failed verifyAllWhenMocksCalled: 1 mock(s) not called',
+      )
+    }
+    expect(error).toBeDefined()
+  })
 })