useLazyQuery: fix rules of React violations

.changeset/shaggy-mirrors-judge.md

@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+Avoid usage of useRef in useInternalState to prevent ref access in render.

.changeset/stupid-planes-nail.md

@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+Fix a bug where `useLazyQuery` would not pick up a client change.

src/react/components/__tests__/client/Mutation.test.tsx

@@ -1348,7 +1348,7 @@ describe("General Mutation testing", () => {
                 if (count === 0) {
                   expect(result.called).toEqual(false);
                   expect(result.loading).toEqual(false);
-                  createTodo();
+                  setTimeout(createTodo, 10);
                 } else if (count === 2 && result) {
                   expect(result.data).toEqual(data);
                   setTimeout(() => {
@@ -1358,7 +1358,7 @@ describe("General Mutation testing", () => {
                   });
                 } else if (count === 3) {
                   expect(result.loading).toEqual(false);
-                  createTodo();
+                  setTimeout(createTodo, 10);
                 } else if (count === 5) {
                   expect(result.data).toEqual(data3);
                 }

src/react/hoc/__tests__/mutations/lifecycle.test.tsx

@@ -92,7 +92,9 @@ describe("graphql(mutation) lifecycle", () => {
       class Container extends React.Component<ChildProps<Props>> {
         render() {
           if (this.props.listId !== 2) return null;
-          this.props.mutate!().then(() => resolve());
+          setTimeout(() => {
+            this.props.mutate!().then(() => resolve());
+          });
           return null;
         }
       }

src/react/hooks/useLazyQuery.ts

@@ -9,7 +9,6 @@ import type {
   LazyQueryHookOptions,
   LazyQueryResultTuple,
   NoInfer,
-  QueryResult,
 } from "../types/types.js";
 import { useInternalState } from "./useQuery.js";
 import { useApolloClient } from "./useApolloClient.js";
@@ -95,30 +94,35 @@ export function useLazyQuery<
     useQueryResult.observable.options.initialFetchPolicy ||
     internalState.getDefaultFetchPolicy();
 
-  const result: QueryResult<TData, TVariables> = Object.assign(useQueryResult, {
-    called: !!execOptionsRef.current,
-  });
-
+  const { forceUpdateState, obsQueryFields } = internalState;
   // We use useMemo here to make sure the eager methods have a stable identity.
   const eagerMethods = React.useMemo(() => {
     const eagerMethods: Record<string, any> = {};
     for (const key of EAGER_METHODS) {
-      const method = result[key];
+      const method = obsQueryFields[key];
       eagerMethods[key] = function () {
         if (!execOptionsRef.current) {
           execOptionsRef.current = Object.create(null);
           // Only the first time populating execOptionsRef.current matters here.
-          internalState.forceUpdateState();
+          forceUpdateState();
         }
         // @ts-expect-error this is just too generic to type
         return method.apply(this, arguments);
       };
     }
 
     return eagerMethods;
-  }, []);
-
-  Object.assign(result, eagerMethods);
+  }, [forceUpdateState, obsQueryFields]);
+
+  const called = !!execOptionsRef.current;
+  const result = React.useMemo(
+    () => ({
+      ...useQueryResult,
+      ...eagerMethods,
+      called,
+    }),
+    [useQueryResult, eagerMethods, called]
+  );
 
   const execute = React.useCallback<LazyQueryResultTuple<TData, TVariables>[0]>(
     (executeOptions) => {
@@ -147,7 +151,7 @@ export function useLazyQuery<
 
       return promise;
     },
-    []
+    [eagerMethods, initialFetchPolicy, internalState]
   );
 
   return [execute, result];

src/react/hooks/useQuery.ts

@@ -109,23 +109,30 @@ export function useInternalState<TData, TVariables extends OperationVariables>(
   client: ApolloClient<any>,
   query: DocumentNode | TypedDocumentNode<TData, TVariables>
 ): InternalState<TData, TVariables> {
-  const stateRef = React.useRef<InternalState<TData, TVariables>>();
-  if (
-    !stateRef.current ||
-    client !== stateRef.current.client ||
-    query !== stateRef.current.query
-  ) {
-    stateRef.current = new InternalState(client, query, stateRef.current);
-  }
-  const state = stateRef.current;
-
   // By default, InternalState.prototype.forceUpdate is an empty function, but
   // we replace it here (before anyone has had a chance to see this state yet)
   // with a function that unconditionally forces an update, using the latest
-  // setTick function. Updating this state by calling state.forceUpdate is the
-  // only way we trigger React component updates (no other useState calls within
-  // the InternalState class).
-  state.forceUpdateState = React.useReducer((tick) => tick + 1, 0)[1];
+  // setTick function. Updating this state by calling state.forceUpdate or the
+  // uSES notification callback are the only way we trigger React component updates.
+  const forceUpdateState = React.useReducer((tick) => tick + 1, 0)[1];
+
+  function createInternalState(previous?: InternalState<TData, TVariables>) {
+    return Object.assign(new InternalState(client, query, previous), {
+      forceUpdateState,
+    });
+  }
+
+  let [state, updateState] = React.useState(createInternalState);
+
+  if (client !== state.client || query !== state.query) {
+    // If the client or query have changed, we need to create a new InternalState.
+    // This will trigger a re-render with the new state, but it will also continue
+    // to run the current render function to completion.
+    // Since we sometimes trigger some side-effects in the render function, we
+    // re-assign `state` to the new state to ensure that those side-effects are
+    // triggered with the new state.
+    updateState((state = createInternalState(state)));
+  }
 
   return state;
 }
@@ -511,7 +518,7 @@ class InternalState<TData, TVariables extends OperationVariables> {
   private onError(error: ApolloError) {}
 
   private observable!: ObservableQuery<TData, TVariables>;
-  private obsQueryFields!: Omit<
+  public obsQueryFields!: Omit<
     ObservableQueryFields<TData, TVariables>,
     "variables"
   >;