feat(zones): fetch zones using react-query MAASENG-3404 (#5485)

- Enable React Query for managing zone-related data fetching and caching
- Add `WebSocketEndpoints` detailing allowed WebSocket endpoint models and methods
- Remove unused imports and code related to fetching zones
- Refactor code to use new `useZoneById` hook for fetching zones by ID
- Update testing utilities to support React Query and WebSocket testing
- Add new helper functions in testing/utils for setting up initial state and query data
- Modify `renderWithBrowserRouter` to return `store` and `queryClient` for more concise tests

src/app/Routes.test.tsx

@@ -109,7 +109,9 @@ const routes: { title: string; path: string }[] = [
 describe("Routes", () => {
   let state: RootState;
   let scrollToSpy: Mock;
-
+  const queryData = {
+    zones: [factory.zone({ id: 1, name: "test-zone" })],
+  };
   beforeEach(() => {
     state = factory.rootState({
       user: factory.userState({
@@ -146,14 +148,7 @@ describe("Routes", () => {
           }),
         ],
       }),
-      zone: factory.zoneState({
-        items: [
-          factory.zone({
-            id: 1,
-            name: "test-zone",
-          }),
-        ],
-      }),
+      zone: factory.zoneState({}),
     });
     scrollToSpy = vi.fn();
     global.scrollTo = scrollToSpy;
@@ -168,6 +163,7 @@ describe("Routes", () => {
       renderWithBrowserRouter(<Routes />, {
         route: path,
         state,
+        queryData,
         routePattern: "/*",
       });
       await waitFor(() => expect(document.title).toBe(`${title} | MAAS`), {

src/app/api/base.ts

@@ -18,7 +18,7 @@ export const handleErrors = (response: Response) => {
 };
 
 type ApiEndpoint = typeof API_ENDPOINTS;
-type ApiEndpointKey = keyof ApiEndpoint;
+export type ApiEndpointKey = keyof ApiEndpoint;
 type ApiUrl = `${typeof SERVICE_API}${ApiEndpoint[ApiEndpointKey]}`;
 
 export const getFullApiUrl = (endpoint: ApiEndpointKey): ApiUrl =>

src/app/api/query-client.ts

@@ -13,6 +13,9 @@ type QueryKeySubcategories<T extends QueryKeyCategories> = keyof QueryKeys[T];
 export type QueryKey =
   QueryKeys[QueryKeyCategories][QueryKeySubcategories<QueryKeyCategories>];
 
+// first element of the queryKeys array
+export type QueryModel = QueryKey[number];
+
 export const defaultQueryOptions = {
   staleTime: 5 * 60 * 1000, // 5 minutes
   cacheTime: 15 * 60 * 1000, // 15 minutes

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

@@ -61,5 +61,6 @@ it("returns the result of useQuery", () => {
   const { result } = renderHookWithMockStore(() =>
     useWebsocketAwareQuery(mockQueryKey, mockQueryFn)
   );
+  expect(result.current).not.toBeNull();
   expect(result.current).toEqual({ data: "testData", isLoading: false });
 });

src/app/api/query/base.ts

@@ -1,12 +1,44 @@
-import { useEffect } from "react";
+import { useEffect, useCallback, useContext } from "react";
 
 import type { QueryFunction, UseQueryOptions } from "@tanstack/react-query";
 import { useQuery, useQueryClient } from "@tanstack/react-query";
 import { useSelector } from "react-redux";
 
 import type { QueryKey } from "@/app/api/query-client";
+import { WebSocketContext } from "@/app/base/websocket-context";
 import statusSelectors from "@/app/store/status/selectors";
+import type { WebSocketEndpointModel } from "@/websocket-client";
+import { WebSocketMessageType } from "@/websocket-client";
 
+export const useWebSocket = () => {
+  const websocketClient = useContext(WebSocketContext);
+
+  if (!websocketClient) {
+    throw new Error("useWebSocket must be used within a WebSocketProvider");
+  }
+
+  const subscribe = useCallback(
+    (callback: (msg: any) => void) => {
+      if (!websocketClient.rws) return;
+
+      const messageHandler = (messageEvent: MessageEvent) => {
+        const data = JSON.parse(messageEvent.data);
+        if (data.type === WebSocketMessageType.NOTIFY) callback(data);
+      };
+      websocketClient.rws.addEventListener("message", messageHandler);
+      return () =>
+        websocketClient.rws?.removeEventListener("message", messageHandler);
+    },
+    [websocketClient]
+  );
+
+  return { subscribe };
+};
+
+const wsToQueryKeyMapping: Partial<Record<WebSocketEndpointModel, string>> = {
+  zone: "zones",
+  // Add more mappings as needed
+} as const;
 export function useWebsocketAwareQuery<
   TQueryFnData = unknown,
   TError = unknown,
@@ -21,11 +53,27 @@ export function useWebsocketAwareQuery<
 ) {
   const queryClient = useQueryClient();
   const connectedCount = useSelector(statusSelectors.connectedCount);
+  const { subscribe } = useWebSocket();
+
+  const queryModelKey = Array.isArray(queryKey) ? queryKey[0] : "";
 
   useEffect(() => {
     queryClient.invalidateQueries();
   }, [connectedCount, queryClient, queryKey]);
 
+  useEffect(() => {
+    return subscribe(
+      ({ name: model }: { action: string; name: WebSocketEndpointModel }) => {
+        const mappedKey = wsToQueryKeyMapping[model];
+        const modelQueryKey = queryKey[0];
+
+        if (mappedKey && mappedKey === modelQueryKey) {
+          queryClient.invalidateQueries({ queryKey });
+        }
+      }
+    );
+  }, [queryClient, subscribe, queryModelKey, queryKey]);
+
   return useQuery<TQueryFnData, TError, TData>({
     queryKey,
     queryFn,

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

@@ -1,46 +1,50 @@
-import type { UseQueryResult } from "@tanstack/react-query";
+import { selectItemsCount, selectById } from "./utils";
 
-import { useItemsCount } from "./utils";
+describe("selectItemsCount", () => {
+  it("should return 0 for undefined input", () => {
+    const count = selectItemsCount()(undefined);
+    expect(count).toBe(0);
+  });
 
-import { renderHook } from "@/testing/utils";
+  it("should return the correct count for a non-empty array", () => {
+    const data = [1, 2, 3, 4, 5];
+    const count = selectItemsCount()(data);
+    expect(count).toBe(5);
+  });
 
-it("should return 0 when data is undefined", () => {
-  const mockUseItems = vi.fn(
-    () => ({ data: undefined }) as UseQueryResult<any[], unknown>
-  );
-  const { result } = renderHook(() => useItemsCount(mockUseItems));
-  expect(result.current).toBe(0);
+  it("should return 0 for an empty array", () => {
+    const data: number[] = [];
+    const count = selectItemsCount()(data);
+    expect(count).toBe(0);
+  });
 });
 
-it("should return the correct count when data is available", () => {
-  const mockData = [1, 2, 3, 4, 5];
-  const mockUseItems = vi.fn(
-    () => ({ data: mockData }) as UseQueryResult<number[], unknown>
-  );
-  const { result } = renderHook(() => useItemsCount(mockUseItems));
-  expect(result.current).toBe(5);
-});
+describe("selectById", () => {
+  const testData = [
+    { id: 1, name: "Item 1" },
+    { id: 2, name: "Item 2" },
+    { id: 3, name: "Item 3" },
+    { id: null, name: "Null ID Item" },
+  ];
 
-it("should return 0 when data is an empty array", () => {
-  const mockUseItems = vi.fn();
-  mockUseItems.mockReturnValueOnce({ data: [] } as UseQueryResult<[], unknown>);
-  const { result } = renderHook(() => useItemsCount(mockUseItems));
-  expect(result.current).toBe(0);
-});
+  it("should return the correct item when given a valid ID", () => {
+    const item = selectById(2)(testData);
+    expect(item).toEqual({ id: 2, name: "Item 2" });
+  });
+
+  it("should return null when given an ID that does not exist", () => {
+    const item = selectById(4)(testData);
+    expect(item).toBeNull();
+  });
+
+  it("should return the correct item when given a null ID", () => {
+    const item = selectById(null)(testData);
+    expect(item).toEqual({ id: null, name: "Null ID Item" });
+  });
 
-it("should update count when data changes", () => {
-  const mockUseItems = vi.fn();
-  mockUseItems.mockReturnValueOnce({ data: [1, 2, 3] } as UseQueryResult<
-    number[],
-    unknown
-  >);
-  const { result, rerender } = renderHook(() => useItemsCount(mockUseItems));
-  expect(result.current).toBe(3);
-
-  mockUseItems.mockReturnValueOnce({ data: [1, 2, 3, 4] } as UseQueryResult<
-    number[],
-    unknown
-  >);
-  rerender();
-  expect(result.current).toBe(4);
+  it("should return null when given a null ID and no matching item exists", () => {
+    const dataWithoutNullId = testData.filter((item) => item.id !== null);
+    const item = selectById(null)(dataWithoutNullId);
+    expect(item).toBeNull();
+  });
 });

src/app/api/query/utils.ts

@@ -1,10 +1,20 @@
-import { useMemo } from "react";
-
-import type { UseQueryResult } from "@tanstack/react-query";
-
-type QueryHook<T> = () => UseQueryResult<T[], unknown>;
+/**
+ * Selector function to get the count of items in an array.
+ * @template T
+ * @returns {function(T[] | undefined): number} A function that takes an array of items and returns the count of items.
+ */
+export const selectItemsCount = <T>() => {
+  return (data: T[] | undefined) => data?.length ?? 0;
+};
 
-export const useItemsCount = <T>(useItems: QueryHook<T>) => {
-  const { data } = useItems();
-  return useMemo(() => data?.length ?? 0, [data]);
+/**
+ * Selector function to find an item by its ID.
+ * @template T
+ * @param {number | null} id - The ID of the item to find.
+ * @returns {function(T[]): T | undefined} A function that takes an array of items and returns the item with the specified ID.
+ */
+export const selectById = <T extends { id: number | null }>(
+  id: number | null
+) => {
+  return (data: T[]) => data.find((item) => item.id === id) || null;
 };

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

@@ -1,36 +1,70 @@
-import type { JsonBodyType } from "msw";
+import type { UseQueryResult } from "@tanstack/react-query";
+import { type JsonBodyType } from "msw";
 
-import { useZonesCount } from "./zones";
+import { useZoneCount, useZoneById, useZones } from "./zones";
 
-import { getFullApiUrl } from "@/app/api/base";
 import * as factory from "@/testing/factories";
 import {
   renderHookWithQueryClient,
   setupMockServer,
   waitFor,
 } from "@/testing/utils";
 
-const { server, http, HttpResponse } = setupMockServer();
+const { mockGet } = setupMockServer();
 
-const setupZonesTest = (mockData: JsonBodyType) => {
-  server.use(
-    http.get(getFullApiUrl("zones"), () => HttpResponse.json(mockData))
-  );
-  return renderHookWithQueryClient(() => useZonesCount());
+const setupTest = (
+  hook: () => ReturnType<
+    typeof useZoneCount | typeof useZoneById | typeof useZones
+  >,
+  mockData: JsonBodyType
+) => {
+  mockGet("zones", mockData);
+  return renderHookWithQueryClient(() => hook()) as {
+    result: { current: UseQueryResult<number> };
+  };
 };
 
-it("should return 0 when zones data is undefined", async () => {
-  const { result } = setupZonesTest(null);
-  await waitFor(() => expect(result.current).toBe(0));
+describe("useZones", () => {
+  it("should return zones data when query succeeds", async () => {
+    const mockZones = [factory.zone(), factory.zone()];
+    const { result } = setupTest(useZones, mockZones);
+
+    await waitFor(() => expect(result.current.isSuccess).toBe(true));
+    expect(result.current.data).toEqual(mockZones);
+  });
 });
 
-it("should return the correct count when zones data is available", async () => {
-  const mockZonesData = [factory.zone(), factory.zone(), factory.zone()];
-  const { result } = setupZonesTest(mockZonesData);
-  await waitFor(() => expect(result.current).toBe(3));
+describe("useZoneById", () => {
+  it("should return specific zone when query succeeds", async () => {
+    const mockZones = [factory.zone({ id: 1 }), factory.zone({ id: 2 })];
+    const { result } = setupTest(() => useZoneById(1), mockZones);
+
+    await waitFor(() => expect(result.current.isSuccess).toBe(true));
+    expect(result.current.data).toEqual(mockZones[0]);
+  });
+
+  it("should return null when zone is not found", async () => {
+    const mockZones = [factory.zone({ id: 1 })];
+    const { result } = setupTest(() => useZoneById(2), mockZones);
+
+    await waitFor(() => expect(result.current.isSuccess).toBe(true));
+    expect(result.current.data).toBeNull();
+  });
 });
 
-it("should return 0 when zones data is an empty array", async () => {
-  const { result } = setupZonesTest([]);
-  await waitFor(() => expect(result.current).toBe(0));
+describe("useZoneCount", () => {
+  it("should return correct count when query succeeds", async () => {
+    const mockZones = [factory.zone(), factory.zone(), factory.zone()];
+    const { result } = setupTest(useZoneCount, mockZones);
+
+    await waitFor(() => expect(result.current.isSuccess).toBe(true));
+    expect(result.current.data).toBe(3);
+  });
+
+  it("should return 0 when zones array is empty", async () => {
+    const { result } = setupTest(useZoneCount, []);
+
+    await waitFor(() => expect(result.current.isSuccess).toBe(true));
+    expect(result.current.data).toBe(0);
+  });
 });

src/app/api/query/zones.ts

@@ -1,9 +1,19 @@
+import { selectById } from "./utils";
+
 import { fetchZones } from "@/app/api/endpoints";
 import { useWebsocketAwareQuery } from "@/app/api/query/base";
-import { useItemsCount } from "@/app/api/query/utils";
+import type { Zone, ZonePK } from "@/app/store/zone/types";
 
 export const useZones = () => {
   return useWebsocketAwareQuery(["zones"], fetchZones);
 };
 
-export const useZonesCount = () => useItemsCount(useZones);
+export const useZoneCount = () =>
+  useWebsocketAwareQuery<Zone[], Zone[], number>(["zones"], fetchZones, {
+    select: (data) => data?.length ?? 0,
+  });
+
+export const useZoneById = (id?: ZonePK | null) =>
+  useWebsocketAwareQuery(["zones"], fetchZones, {
+    select: selectById<Zone>(id ?? null),
+  });