Add proxy assignment to retain original expect methods

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
+};

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 => {