Fix React.createContext in SSR

diff --git a/src/getDataFromTree.ts b/src/getDataFromTree.ts
index 8008f09..cc731fb 100755
--- a/src/getDataFromTree.ts
+++ b/src/getDataFromTree.ts
@@ -140,8 +140,14 @@ export function walkTree(
         return;
       }

+      let isProvider = !!(element.type as any)._context;
+      let previousContextValue: any;
+      if (isProvider) {
+        previousContextValue = ((element.type as any)._context as any)._currentValue;
+      }
+
       let child;
-      if ((element.type as any)._context) {
+      if (isProvider) {
         // A provider - sets the context value before rendering children
         ((element.type as any)._context as any)._currentValue = element.props.value;
         child = element.props.children;
@@ -157,6 +163,9 @@ export function walkTree(
           walkTree(child, context, visitor);
         }
       }
+      if (isProvider) {
+        ((element.type as any)._context as any)._currentValue = previousContextValue;
+      }
     } else {
       // A basic string or dom element, just get children
       if (visitor(element, null, context) === false) {
diff --git a/test/server/server.test.tsx b/test/server/server.test.tsx
index 9b2c93c..96bd196 100644
--- a/test/server/server.test.tsx
+++ b/test/server/server.test.tsx
@@ -262,5 +262,32 @@ describe('SSR', () => {
         expect(markup).toMatch(/Tatooine/);
       });
     });
+    it('should work with React.createContext', async () => {
+      let defaultValue = 'default';
+      let Context = React.createContext(defaultValue);
+
+      let providerValue = 'provider';
+
+      expect(
+        await renderToStringWithData(
+          <React.Fragment>
+            <Context.Provider value={providerValue} />
+            <Context.Consumer>{val => val}</Context.Consumer>
+          </React.Fragment>,
+        ),
+      ).toBe('default');
+
+      expect(
+        await renderToStringWithData(
+          <Context.Provider value={providerValue}>
+            <Context.Consumer>{val => val}</Context.Consumer>
+          </Context.Provider>,
+        ),
+      ).toBe(providerValue);
+
+      expect(await renderToStringWithData(<Context.Consumer>{val => val}</Context.Consumer>)).toBe(
+        defaultValue,
+      );
+    });
   });
 });

src/getDataFromTree.ts

@@ -140,8 +140,14 @@ export function walkTree(
         return;
       }
 
+      let isProvider = !!(element.type as any)._context;
+      let previousContextValue: any;
+      if (isProvider) {
+        previousContextValue = ((element.type as any)._context as any)._currentValue;
+      }
+
       let child;
-      if ((element.type as any)._context) {
+      if (isProvider) {
         // A provider - sets the context value before rendering children
         ((element.type as any)._context as any)._currentValue = element.props.value;
         child = element.props.children;
@@ -157,6 +163,9 @@ export function walkTree(
           walkTree(child, context, visitor);
         }
       }
+      if (isProvider) {
+        ((element.type as any)._context as any)._currentValue = previousContextValue;
+      }
     } else {
       // A basic string or dom element, just get children
       if (visitor(element, null, context) === false) {

test/server/server.test.tsx

@@ -262,5 +262,32 @@ describe('SSR', () => {
         expect(markup).toMatch(/Tatooine/);
       });
     });
+    it('should work with React.createContext', async () => {
+      let defaultValue = 'default';
+      let Context = React.createContext(defaultValue);
+
+      let providerValue = 'provider';
+
+      expect(
+        await renderToStringWithData(
+          <React.Fragment>
+            <Context.Provider value={providerValue} />
+            <Context.Consumer>{val => val}</Context.Consumer>
+          </React.Fragment>,
+        ),
+      ).toBe('default');
+
+      expect(
+        await renderToStringWithData(
+          <Context.Provider value={providerValue}>
+            <Context.Consumer>{val => val}</Context.Consumer>
+          </Context.Provider>,
+        ),
+      ).toBe(providerValue);
+
+      expect(await renderToStringWithData(<Context.Consumer>{val => val}</Context.Consumer>)).toBe(
+        defaultValue,
+      );
+    });
   });
 });