diff --git a/site/src/modules/resources/useAgentContainers.test.tsx b/site/src/modules/resources/useAgentContainers.test.tsx index dbdcdf6f21293..363f8d93223c8 100644 --- a/site/src/modules/resources/useAgentContainers.test.tsx +++ b/site/src/modules/resources/useAgentContainers.test.tsx @@ -4,23 +4,19 @@ import type { WorkspaceAgentListContainersResponse } from "api/typesGenerated"; import * as GlobalSnackbar from "components/GlobalSnackbar/utils"; import { http, HttpResponse } from "msw"; import type { FC, PropsWithChildren } from "react"; -import { QueryClient, QueryClientProvider } from "react-query"; +import { act } from "react"; +import { QueryClientProvider } from "react-query"; import { MockWorkspaceAgent, MockWorkspaceAgentDevcontainer, } from "testHelpers/entities"; +import { createTestQueryClient } from "testHelpers/renderHelpers"; import { server } from "testHelpers/server"; import type { OneWayWebSocket } from "utils/OneWayWebSocket"; import { useAgentContainers } from "./useAgentContainers"; const createWrapper = (): FC => { - const queryClient = new QueryClient({ - defaultOptions: { - queries: { - retry: false, - }, - }, - }); + const queryClient = createTestQueryClient(); return ({ children }) => ( {children} ); @@ -111,22 +107,29 @@ describe("useAgentContainers", () => { ), ); - const { unmount } = renderHook( + const { result, unmount } = renderHook( () => useAgentContainers(MockWorkspaceAgent), { wrapper: createWrapper(), }, ); - // Simulate message event with parsing error + // Wait for initial query to complete + await waitFor(() => { + expect(result.current).toEqual([MockWorkspaceAgentDevcontainer]); + }); + + // Now simulate message event with parsing error const messageHandler = mockSocket.addEventListener.mock.calls.find( (call) => call[0] === "message", )?.[1]; if (messageHandler) { - messageHandler({ - parseError: new Error("Parse error"), - parsedMessage: null, + act(() => { + messageHandler({ + parseError: new Error("Parse error"), + parsedMessage: null, + }); }); } @@ -166,20 +169,27 @@ describe("useAgentContainers", () => { ), ); - const { unmount } = renderHook( + const { result, unmount } = renderHook( () => useAgentContainers(MockWorkspaceAgent), { wrapper: createWrapper(), }, ); - // Simulate error event + // Wait for initial query to complete + await waitFor(() => { + expect(result.current).toEqual([MockWorkspaceAgentDevcontainer]); + }); + + // Now simulate error event const errorHandler = mockSocket.addEventListener.mock.calls.find( (call) => call[0] === "error", )?.[1]; if (errorHandler) { - errorHandler(new Error("WebSocket error")); + act(() => { + errorHandler(new Error("WebSocket error")); + }); } await waitFor(() => { @@ -211,4 +221,36 @@ describe("useAgentContainers", () => { watchAgentContainersSpy.mockRestore(); }); + + it("does not establish WebSocket connection when dev container feature is not enabled", async () => { + const watchAgentContainersSpy = jest.spyOn(API, "watchAgentContainers"); + + server.use( + http.get( + `/api/v2/workspaceagents/${MockWorkspaceAgent.id}/containers`, + () => { + return HttpResponse.json( + { message: "Dev Container feature not enabled." }, + { status: 403 }, + ); + }, + ), + ); + + const { result } = renderHook( + () => useAgentContainers(MockWorkspaceAgent), + { + wrapper: createWrapper(), + }, + ); + + // Wait for the query to complete and error to be processed + await waitFor(() => { + expect(result.current).toBeUndefined(); + }); + + expect(watchAgentContainersSpy).not.toHaveBeenCalled(); + + watchAgentContainersSpy.mockRestore(); + }); }); diff --git a/site/src/modules/resources/useAgentContainers.ts b/site/src/modules/resources/useAgentContainers.ts index e2239fe4666f1..8437fbaed6075 100644 --- a/site/src/modules/resources/useAgentContainers.ts +++ b/site/src/modules/resources/useAgentContainers.ts @@ -14,7 +14,11 @@ export function useAgentContainers( ): readonly WorkspaceAgentDevcontainer[] | undefined { const queryClient = useQueryClient(); - const { data: devcontainers } = useQuery({ + const { + data: devcontainers, + error: queryError, + isLoading: queryIsLoading, + } = useQuery({ queryKey: ["agents", agent.id, "containers"], queryFn: () => API.getAgentContainers(agent.id), enabled: agent.status === "connected", @@ -31,7 +35,7 @@ export function useAgentContainers( ); useEffect(() => { - if (agent.status !== "connected") { + if (agent.status !== "connected" || queryIsLoading || queryError) { return; } @@ -57,7 +61,13 @@ export function useAgentContainers( }); return () => socket.close(); - }, [agent.id, agent.status, updateDevcontainersCache]); + }, [ + agent.id, + agent.status, + queryIsLoading, + queryError, + updateDevcontainersCache, + ]); return devcontainers; }