From 8ae82c50656e65626ed0231dd4a770735f0edfb8 Mon Sep 17 00:00:00 2001
From: Jerko Steiner <jerko.steiner@gmail.com>
Date: Wed, 20 Mar 2019 15:28:05 +0500
Subject: [PATCH] Add ability to specify custom JSX in TestUtils

---
 packages/client/src/test-utils/TestUtils.tsx | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/packages/client/src/test-utils/TestUtils.tsx b/packages/client/src/test-utils/TestUtils.tsx
index d44dab5..dc05726 100644
--- a/packages/client/src/test-utils/TestUtils.tsx
+++ b/packages/client/src/test-utils/TestUtils.tsx
@@ -18,6 +18,10 @@ interface IRenderParams<State> {
   state?: DeepPartial<State>
   connector: Connector<any>
   select: IStateSelector<State, any>
+  customJSX?: <Props>(
+    Component: React.ComponentType<Props>,
+    additionalProps: Record<string, Props>,
+  ) => JSX.Element
 }
 
 export class TestUtils {
@@ -51,10 +55,13 @@ export class TestUtils {
     })(params.state)
     const Component = params.connector.connect(params.select)
 
-    const render = (additionalProps: {[key: string]: any} = {}) => {
+    const render = (additionalProps: Record<string, any> = {}) => {
+      const jsx = params.customJSX
+        ? params.customJSX(Component, additionalProps)
+        : <Component {...additionalProps} />
       return this.render(
         <Provider store={store}>
-          <Component {...additionalProps} />
+          {jsx}
         </Provider>,
       )
     }