From 795ce09a70ba0045cc9851b38dd1ead81196d17d Mon Sep 17 00:00:00 2001
From: Alex Krolick <alexander.krolick@gmail.com>
Date: Sat, 16 Feb 2019 14:12:33 -0800
Subject: [PATCH 1/4] feat(render): add wrapper component option

---
 src/__tests__/render.js | 21 ++++++++++++++++++++
 src/index.js            | 44 +++++++++++++++++++++++------------------
 2 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/src/__tests__/render.js b/src/__tests__/render.js
index 5ee0dc6f..0ffe1322 100644
--- a/src/__tests__/render.js
+++ b/src/__tests__/render.js
@@ -90,3 +90,24 @@ it('supports fragments', () => {
   cleanup()
   expect(document.body.innerHTML).toBe('')
 })
+
+test('renders options.wrapper around node', () => {
+  const WrapperComponent = ({children}) => (
+    <div data-testid="wrapper">{children}</div>
+  )
+
+  const {container, getByTestId} = render(<div data-testid="inner" />, {
+    wrapper: WrapperComponent,
+  })
+
+  expect(getByTestId('wrapper')).toBeInTheDocument()
+  expect(container.firstChild).toMatchInlineSnapshot(`
+<div
+  data-testid="wrapper"
+>
+  <div
+    data-testid="inner"
+  />
+</div>
+`)
+})
diff --git a/src/index.js b/src/index.js
index 31f99339..c0fe6b90 100644
--- a/src/index.js
+++ b/src/index.js
@@ -11,7 +11,13 @@ const mountedContainers = new Set()
 
 function render(
   ui,
-  {container, baseElement = container, queries, hydrate = false} = {},
+  {
+    container,
+    baseElement = container,
+    queries,
+    hydrate = false,
+    wrapper: WrapperComponent,
+  } = {},
 ) {
   if (!container) {
     // default to document.body instead of documentElement to avoid output of potentially-large
@@ -25,13 +31,18 @@ function render(
   // they're passing us a custom container or not.
   mountedContainers.add(container)
 
+  const wrapUiIfNeeded = innerElement =>
+    WrapperComponent
+      ? React.createElement(WrapperComponent, null, innerElement)
+      : innerElement
+
   if (hydrate) {
     act(() => {
-      ReactDOM.hydrate(ui, container)
+      ReactDOM.hydrate(wrapUiIfNeeded(ui), container)
     })
   } else {
     act(() => {
-      ReactDOM.render(ui, container)
+      ReactDOM.render(wrapUiIfNeeded(ui), container)
     })
   }
   return {
@@ -41,7 +52,7 @@ function render(
     debug: (el = baseElement) => console.log(prettyDOM(el)),
     unmount: () => ReactDOM.unmountComponentAtNode(container),
     rerender: rerenderUi => {
-      render(rerenderUi, {container, baseElement})
+      render(wrapUiIfNeeded(rerenderUi), {container, baseElement})
       // Intentionally do not return anything to avoid unnecessarily complicating the API.
       // folks can use all the same utilities we return in the first place that are bound to the container
     },
@@ -70,25 +81,20 @@ function testHook(callback, options = {}) {
   const result = {
     current: null,
   }
-  const toRender = () => {
-    const hookRender = (
-      <TestHook callback={callback}>
-        {res => {
-          result.current = res
-        }}
-      </TestHook>
-    )
-    if (options.wrapper) {
-      return React.createElement(options.wrapper, null, hookRender)
-    }
-    return hookRender
-  }
-  const {unmount, rerender: rerenderComponent} = render(toRender())
+  const toRender = () => (
+    <TestHook callback={callback}>
+      {res => {
+        result.current = res
+      }}
+    </TestHook>
+  )
+
+  const {unmount, rerender: rerenderComponent} = render(toRender(), options)
   return {
     result,
     unmount,
     rerender: () => {
-      rerenderComponent(toRender())
+      rerenderComponent(toRender(), options)
     },
   }
 }

From 1d5717b044eab196d56be0f1f377889619cc8be5 Mon Sep 17 00:00:00 2001
From: Alex Krolick <alexander.krolick@gmail.com>
Date: Sat, 16 Feb 2019 14:25:24 -0800
Subject: [PATCH 2/4] feat(render): update types

---
 package.json       | 3 ++-
 typings/index.d.ts | 7 ++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/package.json b/package.json
index 67bde194..54f531f1 100644
--- a/package.json
+++ b/package.json
@@ -47,7 +47,8 @@
   },
   "devDependencies": {
     "@reach/router": "^1.2.1",
-    "@types/react-dom": "^16.0.9",
+    "@types/react": "^16.8.3",
+    "@types/react-dom": "^16.8.2",
     "axios": "^0.18.0",
     "eslint-import-resolver-jest": "^2.1.1",
     "history": "^4.7.2",
diff --git a/typings/index.d.ts b/typings/index.d.ts
index 16594936..209588cf 100644
--- a/typings/index.d.ts
+++ b/typings/index.d.ts
@@ -25,17 +25,18 @@ export type HookResult<TResult> = {
   unmount: () => boolean
 }
 
-export type HookOptions = {
-  wrapper: React.FunctionComponent
-}
+type WrapperComponent = React.FunctionComponent
 
 export interface RenderOptions<Q extends Queries = typeof queries> {
   container?: HTMLElement
   baseElement?: HTMLElement
   hydrate?: boolean
   queries?: Q
+  wrapper: React.ComponentType
 }
 
+export type HookOptions = RenderOptions
+
 type Omit<T, K extends keyof T> = Pick<T, Exclude<keyof T, K>>
 
 /**

From f63226e2591fed074f5159c44cd7c5b0ae37fdf4 Mon Sep 17 00:00:00 2001
From: Alex Krolick <alexander.krolick@gmail.com>
Date: Sat, 16 Feb 2019 15:29:39 -0800
Subject: [PATCH 3/4] fix(types): remove unreferenced type

---
 typings/index.d.ts | 2 --
 1 file changed, 2 deletions(-)

diff --git a/typings/index.d.ts b/typings/index.d.ts
index 209588cf..e5d122e0 100644
--- a/typings/index.d.ts
+++ b/typings/index.d.ts
@@ -25,8 +25,6 @@ export type HookResult<TResult> = {
   unmount: () => boolean
 }
 
-type WrapperComponent = React.FunctionComponent
-
 export interface RenderOptions<Q extends Queries = typeof queries> {
   container?: HTMLElement
   baseElement?: HTMLElement

From 919c04777cc06e5fbc4cdee83728f4b14cedadf8 Mon Sep 17 00:00:00 2001
From: Alex Krolick <alexkrolick@users.noreply.github.com>
Date: Sat, 16 Feb 2019 17:05:02 -0800
Subject: [PATCH 4/4] fix(render): optional type for wrapper option

---
 typings/index.d.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/typings/index.d.ts b/typings/index.d.ts
index e5d122e0..af09f778 100644
--- a/typings/index.d.ts
+++ b/typings/index.d.ts
@@ -30,7 +30,7 @@ export interface RenderOptions<Q extends Queries = typeof queries> {
   baseElement?: HTMLElement
   hydrate?: boolean
   queries?: Q
-  wrapper: React.ComponentType
+  wrapper?: React.ComponentType
 }
 
 export type HookOptions = RenderOptions