Implement immediate effect and refactor (#250)

* Add useRequest hook to simplify logic

Currently we manually rely on request.key in
useQuery to ensure that we don't trigger a
new effect, when it's not necessary, which
is error-prone.

Instead this hook uses useMemo and checks
request.key against the last request it generated.
If the key hasn't changed it returns the previous
reference.

This ensures reference equality, which removes
any check of request.key in useQuery

* Add useImmediateEffect hook

We want to execute an effect immediately in render
on initial mount for suspense / SSR.

This is a very simple hook that does this by
abstracting effects in three phases:

- Initial mount (render)
- After mount
- Render

So on initial mount it will execute the effect
immediately. Then it will prevent the first
run of useEffect since the effect has already
been called, then the effect will run as usual.

* Fix missing return type in utils/request

* Implement new helpers in useQuery

* Fix useImmediateEffect teardown

* Merge new useEffect and useImmediateEffect

* Adjust useImmediateEffect naming

* Fix initial teardown being ignored in useImmediateEffect

* Update useQuery tests

* Add tests for useImmediateEffect and useRequest

* Add useRequest to useSubscription

* Fix isMounted in hooks

Co-authored-by: Jovi De Croock <[email protected]>

src/hooks/__snapshots__/useQuery.test.tsx.snap

@@ -4,6 +4,6 @@ exports[`on initial useEffect initialises default state 1`] = `
 Object {
   "data": undefined,
   "error": undefined,
-  "fetching": false,
+  "fetching": true,
 }
 `;

src/hooks/useImmediateEffect.test.ts

@@ -0,0 +1,21 @@
+import React from 'react';
+import { renderHook } from 'react-hooks-testing-library';
+import { useImmediateEffect } from './useImmediateEffect';
+
+it('calls effects immediately on mount', () => {
+  const spy = jest.spyOn(React, 'useEffect');
+  const useEffect = jest.fn();
+  const effect = jest.fn();
+
+  spy.mockImplementation(useEffect);
+  renderHook(() => useImmediateEffect(effect, [effect]));
+
+  expect(effect).toHaveBeenCalledTimes(1);
+  expect(effect).toHaveBeenCalledTimes(1);
+
+  expect(useEffect).toHaveBeenCalledWith(expect.any(Function), [
+    expect.any(Function),
+  ]);
+
+  useEffect.mockRestore();
+});

src/hooks/useImmediateEffect.ts

@@ -0,0 +1,37 @@
+import { useRef, useEffect, useCallback } from 'react';
+import { noop } from '../utils';
+
+enum LifecycleState {
+  WillMount = 0,
+  DidMount = 1,
+  Update = 2,
+}
+
+type Effect = () => () => void;
+
+/** This executes an effect immediately on initial render and then treats it as a normal effect */
+export const useImmediateEffect = (
+  effect: Effect,
+  changes: ReadonlyArray<any>
+) => {
+  const teardown = useRef(noop);
+  const state = useRef(LifecycleState.WillMount);
+  const execute = useCallback(effect, changes);
+
+  // On initial render we just execute the effect
+  if (state.current === LifecycleState.WillMount) {
+    state.current = LifecycleState.DidMount;
+    teardown.current = execute();
+  }
+
+  useEffect(() => {
+    // Initially we skip executing the effect since we've already done so on
+    // initial render, then we execute it as usual
+    if (state.current === LifecycleState.Update) {
+      return (teardown.current = execute());
+    } else {
+      state.current = LifecycleState.Update;
+      return teardown.current;
+    }
+  }, [execute]);
+};

src/hooks/useQuery.test.tsx

@@ -20,7 +20,7 @@ jest.mock('../client', () => {
 
 import React, { FC } from 'react';
 import renderer, { act } from 'react-test-renderer';
-import { pipe, onEnd, interval } from 'wonka';
+import { pipe, onStart, onEnd, interval } from 'wonka';
 import { createClient } from '../client';
 import { OperationContext } from '../types';
 import { useQuery, UseQueryArgs, UseQueryState } from './useQuery';
@@ -164,21 +164,23 @@ describe('on change', () => {
 });
 
 describe('on unmount', () => {
+  const start = jest.fn();
   const unsubscribe = jest.fn();
 
   beforeEach(() => {
-    client.executeQuery.mockReturnValueOnce(
+    client.executeQuery.mockReturnValue(
       pipe(
         interval(400),
+        onStart(start),
         onEnd(unsubscribe)
       )
     );
   });
 
-  it('unsubscribe is called', () => {
+  it.only('unsubscribe is called', () => {
     const wrapper = renderer.create(<QueryUser {...props} />);
     wrapper.unmount();
-
+    expect(start).toBeCalledTimes(1);
     expect(unsubscribe).toBeCalledTimes(1);
   });
 });

src/hooks/useQuery.ts

@@ -1,16 +1,11 @@
 import { DocumentNode } from 'graphql';
-import {
-  useCallback,
-  useContext,
-  useEffect,
-  useMemo,
-  useRef,
-  useState,
-} from 'react';
+import { useCallback, useContext, useRef, useState } from 'react';
 import { pipe, subscribe } from 'wonka';
 import { Context } from '../context';
 import { OperationContext, RequestPolicy } from '../types';
-import { CombinedError, createRequest, noop } from '../utils';
+import { CombinedError, noop } from '../utils';
+import { useRequest } from './useRequest';
+import { useImmediateEffect } from './useImmediateEffect';
 
 export interface UseQueryArgs<V> {
   query: string | DocumentNode;
@@ -34,33 +29,25 @@ export const useQuery = <T = any, V = object>(
   args: UseQueryArgs<V>
 ): UseQueryResponse<T> => {
   const isMounted = useRef(true);
-  const unsubscribe = useRef<() => void>(noop);
-
+  const unsubscribe = useRef(noop);
   const client = useContext(Context);
-  const request = useMemo(
-    () => createRequest(args.query, args.variables as any),
-    [args.query, args.variables]
-  );
+
+  // This creates a request which will keep a stable reference
+  // if request.key doesn't change
+  const request = useRequest(args.query, args.variables);
 
   const [state, setState] = useState<UseQueryState<T>>({
     fetching: false,
     error: undefined,
     data: undefined,
   });
 
-  /** Unmount handler */
-  useEffect(
-    () => () => {
-      isMounted.current = false;
-    },
-    []
-  );
-
   const executeQuery = useCallback(
     (opts?: Partial<OperationContext>) => {
       unsubscribe.current();
 
       if (args.pause) {
+        setState(s => ({ ...s, fetching: false }));
         unsubscribe.current = noop;
         return;
       }
@@ -83,15 +70,19 @@ export const useQuery = <T = any, V = object>(
 
       unsubscribe.current = teardown;
     },
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    [request.key, client, args.pause, args.requestPolicy]
+    [request, client, args.pause, args.requestPolicy]
   );
 
-  /** Trigger query on arg change. */
-  useEffect(() => {
+  // Calls executeQuery on initial render immediately, then
+  // treats it as a normal effect
+  useImmediateEffect(() => {
+    isMounted.current = true;
     executeQuery();
 
-    return unsubscribe.current;
+    return () => {
+      isMounted.current = false;
+      unsubscribe.current();
+    };
   }, [executeQuery]);
 
   return [state, executeQuery];

src/hooks/useRequest.test.ts

@@ -0,0 +1,31 @@
+import { renderHook } from 'react-hooks-testing-library';
+import { queryGql } from '../test-utils';
+import { useRequest } from './useRequest';
+
+it('preserves instance of request when key has not changed', () => {
+  let { query, variables } = queryGql;
+
+  const { result, rerender } = renderHook(
+    ({ query, variables }) => useRequest(query, variables),
+    { initialProps: { query, variables } }
+  );
+
+  const resultA = result.current;
+  expect(resultA).toEqual({
+    key: expect.any(Number),
+    query: expect.anything(),
+    variables: variables,
+  });
+
+  variables = { ...variables }; // Change reference
+  rerender({ query, variables });
+
+  const resultB = result.current;
+  expect(resultA).toBe(resultB);
+
+  variables = { ...variables, test: true }; // Change values
+  rerender({ query, variables });
+
+  const resultC = result.current;
+  expect(resultA).not.toBe(resultC);
+});

src/hooks/useRequest.ts

@@ -0,0 +1,23 @@
+import { DocumentNode } from 'graphql';
+import { useRef, useMemo } from 'react';
+import { GraphQLRequest } from '../types';
+import { createRequest } from '../utils';
+
+/** Creates a request from a query and variables but preserves reference equality if the key isn't changing */
+export const useRequest = (
+  query: string | DocumentNode,
+  variables?: any
+): GraphQLRequest => {
+  const prev = useRef<void | GraphQLRequest>(undefined);
+
+  return useMemo(() => {
+    const request = createRequest(query, variables);
+    // We manually ensure reference equality if the key hasn't changed
+    if (prev.current !== undefined && prev.current.key === request.key) {
+      return prev.current;
+    } else {
+      prev.current = request;
+      return request;
+    }
+  }, [query, variables]);
+};

src/hooks/useSubscription.ts

@@ -1,15 +1,9 @@
 import { DocumentNode } from 'graphql';
-import {
-  useCallback,
-  useContext,
-  useEffect,
-  useRef,
-  useState,
-  useMemo,
-} from 'react';
+import { useCallback, useContext, useEffect, useRef, useState } from 'react';
 import { pipe, subscribe } from 'wonka';
 import { Context } from '../context';
-import { CombinedError, createRequest, noop } from '../utils';
+import { CombinedError, noop } from '../utils';
+import { useRequest } from './useRequest';
 
 export interface UseSubscriptionArgs<V> {
   query: DocumentNode | string;
@@ -32,27 +26,18 @@ export const useSubscription = <T = any, R = T, V = object>(
 ): UseSubscriptionResponse<R> => {
   const isMounted = useRef(true);
   const unsubscribe = useRef(noop);
-
   const client = useContext(Context);
-  const request = useMemo(
-    () => createRequest(args.query, args.variables as any),
-    [args.query, args.variables]
-  );
+
+  // This creates a request which will keep a stable reference
+  // if request.key doesn't change
+  const request = useRequest(args.query, args.variables);
 
   const [state, setState] = useState<UseSubscriptionState<R>>({
     fetching: true,
     error: undefined,
     data: undefined,
   });
 
-  /** Unmount handler */
-  useEffect(
-    () => () => {
-      isMounted.current = false;
-    },
-    []
-  );
-
   const executeSubscription = useCallback(() => {
     unsubscribe.current();
 
@@ -74,11 +59,15 @@ export const useSubscription = <T = any, R = T, V = object>(
     unsubscribe.current = teardown;
   }, [client, handler, request]);
 
-  /** Trigger subscription on query change. */
+  // Trigger subscription on query change
   useEffect(() => {
+    isMounted.current = true;
     executeSubscription();
 
-    return unsubscribe.current;
+    return () => {
+      isMounted.current = false;
+      unsubscribe.current();
+    };
   }, [executeSubscription]);
 
   return [state];

src/utils/request.ts

@@ -1,8 +1,12 @@
 import { DocumentNode } from 'graphql';
 import gql from 'graphql-tag';
 import { getKeyForRequest } from './keyForQuery';
+import { GraphQLRequest } from '../types';
 
-export const createRequest = (q: string | DocumentNode, vars?: object) => {
+export const createRequest = (
+  q: string | DocumentNode,
+  vars?: object
+): GraphQLRequest => {
   const query = typeof q === 'string' ? gql([q]) : q;
 
   return {