From 6190463734b92d43dcc974570292ac831c1d1a6c Mon Sep 17 00:00:00 2001
From: Tim Kindberg <tim@vndly.com>
Date: Sat, 6 Feb 2021 19:09:46 -0500
Subject: [PATCH] Support function matchers

---
 src/when.js      | 12 +++++++++++-
 src/when.test.js | 30 ++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/src/when.js b/src/when.js
index 2eed39e..acb63d4 100644
--- a/src/when.js
+++ b/src/when.js
@@ -21,7 +21,17 @@ const checkArgumentMatchers = (expectCall, args) => (match, matcher, i) => {
 
   // Assert the match for better messaging during a failure
   if (expectCall) {
-    expect(arg).toEqual(matcher)
+    if (typeof matcher === 'function') {
+      const isMatch = matcher(arg)
+      const msg = `Failed function matcher within expectCalledWith: ${matcher.name}(${JSON.stringify(arg)}) did not return true\n\n\n...rest of the stack...`
+      assert.equal(isMatch, true, msg)
+    } else {
+      expect(arg).toEqual(matcher)
+    }
+  }
+
+  if (typeof matcher === 'function') {
+    return matcher(arg)
   }
 
   return utils.equals(arg, matcher)
diff --git a/src/when.test.js b/src/when.test.js
index b88c75d..d59178e 100644
--- a/src/when.test.js
+++ b/src/when.test.js
@@ -185,6 +185,36 @@ describe('When', () => {
       expect(fn(1, 'foo', true, 'whatever', undefined, 'oops')).toEqual(undefined)
     })
 
+    it('works with custom function args', () => {
+      const fn = jest.fn()
+
+      const allValuesTrue = (arg) => Object.values(arg).every(Boolean)
+      const numberDivisibleBy3 = (arg) => arg % 3 === 0
+
+      when(fn)
+        .calledWith(allValuesTrue, numberDivisibleBy3)
+        .mockReturnValue('x')
+
+      expect(fn({ foo: true, bar: true }, 9)).toEqual('x')
+      expect(fn({ foo: true, bar: false }, 9)).toEqual(undefined)
+      expect(fn({ foo: true, bar: false }, 13)).toEqual(undefined)
+    })
+
+    it('expects with custom function args', () => {
+      const fn = jest.fn()
+
+      const allValuesTrue = (arg) => Object.values(arg).every(Boolean)
+      const numberDivisibleBy3 = (arg) => arg % 3 === 0
+
+      when(fn)
+        .expectCalledWith(allValuesTrue, numberDivisibleBy3)
+        .mockReturnValue('x')
+
+      expect(fn({ foo: true, bar: true }, 9)).toEqual('x')
+      expect(() => fn({ foo: false, bar: true }, 9)).toThrow(/Failed function matcher within expectCalledWith: allValuesTrue\(\{"foo":false,"bar":true\}\) did not return true/)
+      expect(() => fn({ foo: true, bar: true }, 13)).toThrow(/Failed function matcher within expectCalledWith: numberDivisibleBy3\(13\) did not return true/)
+    })
+
     it('supports compound when declarations', () => {
       const fn = jest.fn()