From 264a3c7becd6736fe0a2b527777e9624c99f509e Mon Sep 17 00:00:00 2001
From: Matt Phillips <matt@mattphillips.io>
Date: Tue, 7 Aug 2018 15:13:18 +0100
Subject: [PATCH] Add proxy assignment to retain original expect methods

---
 src/withMessage.js      |  9 ++++++---
 src/withMessage.test.js | 18 ++++++++++++++++--
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/withMessage.js b/src/withMessage.js
index a6525d9..8280802 100644
--- a/src/withMessage.js
+++ b/src/withMessage.js
@@ -9,9 +9,7 @@ class JestAssertionError extends Error {
   }
 }
 
-export default expect => (actual, customMessage) => {
-  const matchers = expect(actual);
-
+const wrapMatchers = (matchers, customMessage) => {
   return Object.keys(matchers).reduce((acc, key) => {
     const matcher = matchers[key];
     const newMatcher = (...args) => {
@@ -31,3 +29,8 @@ export default expect => (actual, customMessage) => {
     return Object.assign({}, acc, { [key]: newMatcher });
   }, {});
 };
+
+export default expect => {
+  const expectProxy = (actual, customMessage) => wrapMatchers(expect(actual), customMessage); // partially apply expect to get all matchers and wrap them
+  return Object.assign(expectProxy, expect); // clone additional properties on expect
+};
diff --git a/src/withMessage.test.js b/src/withMessage.test.js
index 8d7254a..32eb70f 100644
--- a/src/withMessage.test.js
+++ b/src/withMessage.test.js
@@ -3,13 +3,27 @@ import withMessage from './withMessage';
 describe('withMessage()', () => {
   const ACTUAL = 'ACTUAL';
 
+  test('does not remove additional methods from expect', () => {
+    expect.assertions(3);
+    const toBeMock = jest.fn();
+    const expectMock = jest.fn(() => ({ toBe: toBeMock }));
+    expectMock.extend = 'extend';
+
+    const newExpect = withMessage(expectMock);
+    newExpect(ACTUAL, 'should fail').toBe(1);
+    expect(newExpect.extend).toBe('extend');
+    expect(expectMock).toHaveBeenCalledWith(ACTUAL);
+    expect(toBeMock).toHaveBeenCalledWith(1);
+  });
+
   test('does not throw when matcher passes', () => {
-    expect.assertions(1);
+    expect.assertions(2);
     const toBeMock = jest.fn();
     const expectMock = jest.fn(() => ({ toBe: toBeMock }));
 
     withMessage(expectMock)(ACTUAL, 'should fail').toBe(1);
-    expect(1).toBe(1);
+    expect(expectMock).toHaveBeenCalledWith(ACTUAL);
+    expect(toBeMock).toHaveBeenCalledWith(1);
   });
 
   test.each([undefined, ''])('throws original error when given message: %s', message => {