From fdbe4be75f0ce34ab8067489eda9fc392aaeab25 Mon Sep 17 00:00:00 2001
From: Jerko Steiner <jerko.steiner@gmail.com>
Date: Wed, 11 Sep 2019 10:13:34 +0700
Subject: [PATCH] Return waitForActions after rendering

---
 .../client/src/team/TeamConnector.test.tsx    | 11 +++++-----
 packages/client/src/test-utils/TestUtils.tsx  | 20 ++++++++++++-------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/packages/client/src/team/TeamConnector.test.tsx b/packages/client/src/team/TeamConnector.test.tsx
index b1dcd78..f2bf940 100644
--- a/packages/client/src/team/TeamConnector.test.tsx
+++ b/packages/client/src/team/TeamConnector.test.tsx
@@ -69,8 +69,7 @@ describe('TeamConnector', () => {
   }
 
   it('it fetches user teams on render', async () => {
-    const {waitForActions, render} = createTestProvider()
-    const {node} = render({
+    const {node, waitForActions} = createTestProvider().render({
       history,
       location: {} as any,
       match: {} as any,
@@ -94,8 +93,8 @@ describe('TeamConnector', () => {
         userTeams: [],
       }
       teamClientMock.create.mockResolvedValue(newTeam)
-      const {render, store, waitForActions} = createTestProvider()
-      const {node} = render({
+      const {render, store} = createTestProvider()
+      const {waitForActions, node} = render({
         history,
         location: {} as any,
         match: {} as any,
@@ -114,8 +113,8 @@ describe('TeamConnector', () => {
     it('displays an error', async () => {
       const error = {error: 'An error'}
       teamClientMock.create.mockRejectedValue(new Error('Test Error'))
-      const {render, waitForActions} = createTestProvider()
-      const {node} = render({
+      const {render} = createTestProvider()
+      const {node, waitForActions} = render({
         history,
         location: {} as any,
         match: {} as any,
diff --git a/packages/client/src/test-utils/TestUtils.tsx b/packages/client/src/test-utils/TestUtils.tsx
index 33d84e0..116baef 100644
--- a/packages/client/src/test-utils/TestUtils.tsx
+++ b/packages/client/src/test-utils/TestUtils.tsx
@@ -92,11 +92,17 @@ export class TestUtils {
         const jsx = createJSX
           ? createJSX(Component, props)
           : <Component {...props} />
-        return this.render(
+        const result = this.render(
           <Provider store={store}>
             {jsx}
           </Provider>,
         )
+        return {
+          ...result,
+          async waitForActions(timeout = 2000) {
+            await waitMiddleware.waitForRecorded(recorder)
+          },
+        }
       }
 
       const withJSX = (localCreateJSX: CreateJSX) => {
@@ -106,14 +112,11 @@ export class TestUtils {
 
       const self: ISelf<
         Props, typeof store, typeof Component, CreateJSX
-        > = {
+      > = {
         render,
         store,
         Component,
         withJSX,
-        async waitForActions() {
-          await waitMiddleware.waitForRecorded(recorder, 2000)
-        },
       }
 
       return self
@@ -124,10 +127,13 @@ export class TestUtils {
 }
 
 interface ISelf<Props, Store, Component, CreateJSX> {
-  render: (props: Props) => ReturnType<TestUtils['render']>
+  render: (props: Props) => {
+    component: React.Component<any>
+    node: Element
+    waitForActions(timeout?: number): Promise<void>
+  }
   store: Store
   Component: Component
   withJSX: (localCreateJSX: CreateJSX)
   => ISelf<Props, Store, Component, CreateJSX>
-  waitForActions(): Promise<void>
 }