fix(zones): unnecessary query invalidations and API calls (#5518)

- refactor: invalidate queries only when connectedCount changes
- refactor: add zoneKeys

src/app/api/query/base.test.ts

@@ -34,6 +34,31 @@ it("calls useQuery with correct parameters", () => {
   });
 });
 
+it("skips query invalidation when connectedCount is unchanged", () => {
+  const initialState = rootState({
+    status: statusState({ connectedCount: 0 }),
+  });
+  const { rerender } = renderHookWithMockStore(
+    () => useWebsocketAwareQuery(mockQueryKey, mockQueryFn),
+    { initialState }
+  );
+
+  const mockInvalidateQueries = vi.fn();
+  const mockQueryClient: Partial<reactQuery.QueryClient> = {
+    invalidateQueries: mockInvalidateQueries,
+  };
+  vi.mocked(reactQuery.useQueryClient).mockReturnValue(
+    mockQueryClient as reactQuery.QueryClient
+  );
+
+  rerender(() => useWebsocketAwareQuery(mockQueryKey, mockQueryFn), {
+    state: rootState({
+      status: statusState({ connectedCount: 0 }),
+    }),
+  });
+  expect(mockInvalidateQueries).not.toHaveBeenCalled();
+});
+
 it("invalidates queries when connectedCount changes", () => {
   const initialState = rootState({
     status: statusState({ connectedCount: 0 }),
@@ -51,8 +76,10 @@ it("invalidates queries when connectedCount changes", () => {
     mockQueryClient as reactQuery.QueryClient
   );
 
-  rerender({
-    initialState: rootState({ status: statusState({ connectedCount: 1 }) }),
+  rerender(() => useWebsocketAwareQuery(mockQueryKey, mockQueryFn), {
+    state: rootState({
+      status: statusState({ connectedCount: 1 }),
+    }),
   });
   expect(mockInvalidateQueries).toHaveBeenCalled();
 });

src/app/api/query/base.ts

@@ -1,5 +1,6 @@
 import { useEffect, useCallback, useContext } from "react";
 
+import { usePrevious } from "@canonical/react-components";
 import type { QueryFunction, UseQueryOptions } from "@tanstack/react-query";
 import { useQuery, useQueryClient } from "@tanstack/react-query";
 import { useSelector } from "react-redux";
@@ -56,10 +57,13 @@ export function useWebsocketAwareQuery<
   const { subscribe } = useWebSocket();
 
   const queryModelKey = Array.isArray(queryKey) ? queryKey[0] : "";
+  const previousConnectedCount = usePrevious(connectedCount);
 
   useEffect(() => {
-    queryClient.invalidateQueries();
-  }, [connectedCount, queryClient, queryKey]);
+    if (connectedCount !== previousConnectedCount) {
+      queryClient.invalidateQueries({ queryKey });
+    }
+  }, [connectedCount, previousConnectedCount, queryClient, queryKey]);
 
   useEffect(() => {
     return subscribe(

src/app/api/query/zones.ts

@@ -4,8 +4,12 @@ import { fetchZones } from "@/app/api/endpoints";
 import { useWebsocketAwareQuery } from "@/app/api/query/base";
 import type { Zone, ZonePK } from "@/app/store/zone/types";
 
+const zoneKeys = {
+  list: ["zones"] as const,
+};
+
 export const useZones = () => {
-  return useWebsocketAwareQuery(["zones"], fetchZones);
+  return useWebsocketAwareQuery(zoneKeys.list, fetchZones);
 };
 
 export const useZoneCount = () =>
@@ -14,6 +18,6 @@ export const useZoneCount = () =>
   });
 
 export const useZoneById = (id?: ZonePK | null) =>
-  useWebsocketAwareQuery(["zones"], fetchZones, {
+  useWebsocketAwareQuery(zoneKeys.list, fetchZones, {
     select: selectById<Zone>(id ?? null),
   });

src/testing/utils.tsx

@@ -413,22 +413,35 @@ export const expectTooltipOnHover = async (
   }
 };
 
-const generateWrapper =
-  (store = configureStore()(rootStateFactory())) =>
-  ({ children }: { children: ReactNode }) => (
+type Hook = Parameters<typeof renderHook>[0];
+export const renderHookWithMockStore = (
+  hook: Hook,
+  options?: { initialState?: RootState }
+) => {
+  let store = configureStore()(options?.initialState || rootStateFactory());
+  const wrapper = ({ children }: { children: ReactNode }) => (
     <WebSocketProvider>
       <Provider store={store}>{children}</Provider>
     </WebSocketProvider>
   );
 
-type Hook = Parameters<typeof renderHook>[0];
+  const result = renderHook(hook, { wrapper });
 
-export const renderHookWithMockStore = (
-  hook: Hook,
-  options?: { initialState?: RootState }
-) => {
-  const store = configureStore()(options?.initialState || rootStateFactory());
-  return renderHook(hook, { wrapper: generateWrapper(store) });
+  const customRerender = (
+    newHook?: Hook,
+    { state: newState }: { state?: Partial<RootState> } = {}
+  ) => {
+    if (newState) {
+      store = configureStore()({ ...newState });
+    }
+    result.rerender(newHook);
+  };
+
+  return {
+    ...result,
+    rerender: customRerender,
+    store,
+  };
 };
 
 export const renderHookWithQueryClient = (hook: Hook) => {