Fix broken SSR on master (#310)

* fix broken ssr

* update tests

* update tests

* remove ssr abstraction

* don't call useLayoutEffect in SSR

* update SSR check

* ssr....

src/hooks/useImmediateState.test.tsx

@@ -0,0 +1,47 @@
+import React from 'react';
+import renderer from 'react-test-renderer';
+import { useImmediateState } from './useImmediateState';
+
+const setStateMock = jest.fn();
+jest.spyOn(React, 'useState').mockImplementation(arg => [arg, setStateMock]);
+
+const initialState = { someObject: 1234 };
+const updateState = { someObject: 5678 };
+let state;
+let setState;
+
+const Fixture = ({ update }: { update?: boolean }) => {
+  const [a, set] = useImmediateState<object>(initialState);
+
+  if (update) {
+    set(updateState);
+  }
+
+  state = a;
+  setState = set;
+
+  return null;
+};
+
+beforeEach(jest.clearAllMocks);
+
+describe('on initial mount', () => {
+  it('sets initial state', () => {
+    renderer.create(<Fixture />);
+    expect(state).toEqual(initialState);
+  });
+
+  it('only mutates on setState call', () => {
+    renderer.create(<Fixture update={true} />);
+    expect(setStateMock).toBeCalledTimes(0);
+    expect(state).toEqual(updateState);
+  });
+});
+
+describe('on later mounts', () => {
+  it('sets state via setState', () => {
+    renderer.create(<Fixture />);
+    setState(updateState);
+    expect(setStateMock).toBeCalledTimes(1);
+  });
+});

src/hooks/useImmediateState.ts

@@ -1,39 +1,39 @@
-import { useRef, useEffect, useState, useCallback } from 'react';
+import { useRef, useState, useCallback, useLayoutEffect } from 'react';
+import { isSSR } from '../utils';
 
 type SetStateAction<S> = S | ((prevState: S) => S);
 type SetState<S> = (action: SetStateAction<S>) => void;
 
-/** This is a drop-in replacement for useState, limited to object-based state. During initial mount it will mutably update the state, instead of scheduling a React update using setState */
+/**
+ * This is a drop-in replacement for useState, limited to object-based state.
+ * During initial mount it will mutably update the state, instead of scheduling
+ * a React update using setState
+ */
 export const useImmediateState = <S extends {}>(init: S): [S, SetState<S>] => {
   const isMounted = useRef(false);
   const initialState = useRef<S>({ ...init });
   const [state, setState] = useState<S>(initialState.current);
 
   // This wraps setState and updates the state mutably on initial mount
-  // It also prevents setting the state when the component is unmounted
   const updateState: SetState<S> = useCallback((action: SetStateAction<S>) => {
-    if (isMounted.current) {
-      setState(action);
-    } else if (typeof action === 'function') {
-      const update = (action as any)(initialState.current);
-      Object.assign(initialState.current, update);
-    } else {
-      // There are scenario's where we are in-between mounts.
-      // The reason we need both is because we COULD still be mounting
-      // and we could be in the process of being mounted.
-      // Scenario 1 needs the initialState.current mutation.
-      // Scenario 2 needs the setState.
-      // See https://github.com/FormidableLabs/urql/issues/287
-      setState(() => Object.assign(initialState.current, action as any));
+    if (!isMounted.current) {
+      const newValue =
+        typeof action === 'function'
+          ? (action as (arg: S) => S)(initialState.current)
+          : action;
+      return Object.assign(initialState.current, newValue);
     }
-  }, []);
 
-  useEffect(() => {
-    isMounted.current = true;
-    return () => {
-      isMounted.current = false;
-    };
+    setState(action);
   }, []);
 
+  !isSSR && // eslint-disable-next-line react-hooks/rules-of-hooks
+    useLayoutEffect(() => {
+      isMounted.current = true;
+      return () => {
+        isMounted.current = false;
+      };
+    }, []);
+
   return [state, updateState];
 };

src/utils/index.ts

@@ -1,6 +1,7 @@
 export { CombinedError } from './error';
 export { getKeyForRequest } from './keyForQuery';
 export { createRequest } from './request';
+export * from './ssr';
 export { formatDocument, collectTypesFromResponse } from './typenames';
 export { toSuspenseSource } from './toSuspenseSource';
 

src/utils/ssr.ts

@@ -0,0 +1,2 @@
+export const isSSR =
+  typeof window === 'undefined' || !('HTMLElement' in window);