Replace ref with setState, batch updates

Author: msutkowski

src/react-hooks/buildHooks.ts

@@ -60,7 +60,7 @@ export type UseLazyQuery<D extends QueryDefinition<any, any, any, any>> = <R = U
 
 export type UseLazyQuerySubscription<D extends QueryDefinition<any, any, any, any>> = (
   options?: SubscriptionOptions
-) => [(args: QueryArgFrom<D>) => void, undefined | React.MutableRefObject<QueryActionCreatorResult<any>>['current']];
+) => [(args: QueryArgFrom<D>) => void, undefined | QueryActionCreatorResult<any>];
 
 export type QueryStateSelector<R, D extends QueryDefinition<any, any, any, any>> = (
   state: QueryResultSelectorResult<D>,
@@ -252,57 +252,62 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
       refetchOnFocus,
       pollingInterval = 0,
     } = {}) => {
-      const forceUpdate = useReducer(() => ({}), {})[1];
       const { initiate } = api.endpoints[name] as ApiEndpointQuery<
         QueryDefinition<any, any, any, any, any>,
         Definitions
       >;
       const dispatch = useDispatch<ThunkDispatch<any, any, AnyAction>>();
 
-      const promiseRef = useRef<QueryActionCreatorResult<any>>();
-      const lastPromise = promiseRef.current;
+      const [promiseRef, setPromiseRef] = useState<QueryActionCreatorResult<any>>();
+      const unmountRef = useRef<QueryActionCreatorResult<any> | undefined>();
 
       useEffect(() => {
         const subscriptionOptions = {
           refetchOnReconnect,
           refetchOnFocus,
           pollingInterval,
         };
-        if (lastPromise && !shallowEqual(subscriptionOptions, lastPromise.subscriptionOptions)) {
-          lastPromise?.updateSubscriptionOptions(subscriptionOptions);
-          lastPromise.subscriptionOptions = subscriptionOptions;
+        if (promiseRef && !shallowEqual(subscriptionOptions, promiseRef.subscriptionOptions)) {
+          promiseRef?.updateSubscriptionOptions(subscriptionOptions);
+          setPromiseRef((prev) => ({ ...prev!, subscriptionOptions }));
         }
-      }, [lastPromise, refetchOnFocus, refetchOnReconnect, pollingInterval]);
+      }, [promiseRef, refetchOnFocus, refetchOnReconnect, pollingInterval]);
+
+      useEffect(() => {
+        unmountRef.current = promiseRef;
+      }, [promiseRef]);
 
       useEffect(() => {
         return () => {
-          promiseRef.current?.unsubscribe();
-          promiseRef.current = undefined;
+          unmountRef?.current?.unsubscribe();
         };
       }, []);
 
       const trigger = useCallback(
         function (arg: any) {
-          lastPromise?.unsubscribe();
-
-          const subscriptionOptions = {
-            refetchOnReconnect,
-            refetchOnFocus,
-            pollingInterval,
-          };
-
-          promiseRef.current = dispatch(
-            initiate(arg, {
-              subscriptionOptions,
-              forceRefetch: true,
-            })
-          );
-          forceUpdate();
+          batch(() => {
+            promiseRef?.unsubscribe();
+
+            const subscriptionOptions = {
+              refetchOnReconnect,
+              refetchOnFocus,
+              pollingInterval,
+            };
+
+            const promise = dispatch(
+              initiate(arg, {
+                subscriptionOptions,
+                forceRefetch: true,
+              })
+            );
+
+            setPromiseRef(promise);
+          });
         },
-        [dispatch, initiate, lastPromise, pollingInterval, refetchOnFocus, refetchOnReconnect]
+        [dispatch, initiate, promiseRef, pollingInterval, refetchOnFocus, refetchOnReconnect]
       );
 
-      return useMemo(() => [trigger, lastPromise], [trigger, lastPromise]);
+      return useMemo(() => [trigger, promiseRef], [trigger, promiseRef]);
     };
 
     const useQueryState: UseQueryState<any> = (

test/buildHooks.test.tsx

@@ -488,7 +488,7 @@ describe('hooks tests', () => {
         );
       }
 
-      render(<User />, { wrapper: storeRef.wrapper });
+      const { unmount } = render(<User />, { wrapper: storeRef.wrapper });
 
       await waitFor(() => expect(screen.getByTestId('isUninitialized').textContent).toBe('true'));
       await waitFor(() => expect(data).toBeUndefined());
@@ -523,6 +523,13 @@ describe('hooks tests', () => {
       expect(storeRef.store.getState().actions.filter(api.internalActions.unsubscribeQueryResult.match)).toHaveLength(
         3
       );
+
+      unmount();
+
+      // We unsubscribe after the component unmounts
+      expect(storeRef.store.getState().actions.filter(api.internalActions.unsubscribeQueryResult.match)).toHaveLength(
+        4
+      );
     });
   });