Add useAgentLogs test

coder · BrunoQuaresma · Jun 7, 2024 · Jun 3, 2024 · Jun 3, 2024 · Jun 3, 2024
commit 35368ee06de19798e25bca43bb06bb9b5fd2c01a
diff --git a/site/src/api/queries/util.ts b/site/src/api/queries/util.ts
@@ -1,7 +1,7 @@
 import type { UseQueryOptions, QueryKey } from "react-query";
 import type { MetadataState, MetadataValue } from "hooks/useEmbeddedMetadata";
 
-const disabledFetchOptions = {
+export const disabledFetchOptions = {
   cacheTime: Infinity,
   staleTime: Infinity,
   refetchOnMount: false,

diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts
@@ -14,6 +14,7 @@ import type {
   WorkspacesRequest,
   WorkspacesResponse,
 } from "api/typesGenerated";
+import { disabledFetchOptions } from "./util";
 import { workspaceBuildsKey } from "./workspaceBuilds";
 
 export const workspaceByOwnerAndNameKey = (owner: string, name: string) => [
@@ -309,5 +310,6 @@ export const agentLogs = (workspaceId: string, agentId: string) => {
   return {
     queryKey: agentLogsKey(workspaceId, agentId),
     queryFn: () => API.getWorkspaceAgentLogs(agentId),
+    ...disabledFetchOptions,
   };
 };
diff --git a/site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx b/site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx
@@ -0,0 +1,137 @@
+import { act, renderHook, waitFor } from "@testing-library/react";
+import WS from "jest-websocket-mock";
+import { type QueryClient, QueryClientProvider } from "react-query";
+import { API } from "api/api";
+import * as APIModule from "api/api";
+import { agentLogsKey } from "api/queries/workspaces";
+import type {
+  WorkspaceAgentLifecycle,
+  WorkspaceAgentLog,
+} from "api/typesGenerated";
+import { MockWorkspace, MockWorkspaceAgent } from "testHelpers/entities";
+import { createTestQueryClient } from "testHelpers/renderHelpers";
+import { type UseAgentLogsOptions, useAgentLogs } from "./useAgentLogs";
+
+afterEach(() => {
+  WS.clean();
+});
+
+describe("useAgentLogs", () => {
+  it("should not fetch logs if disabled", async () => {
+    const queryClient = createTestQueryClient();
+    const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
+    const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
+    renderUseAgentLogs(queryClient, "ready", { enabled: false });
+    expect(fetchSpy).not.toHaveBeenCalled();
+    expect(wsSpy).not.toHaveBeenCalled();
+  });
+
+  it("should return existing logs without network calls", async () => {
+    const queryClient = createTestQueryClient();
+    queryClient.setQueryData(
+      agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
+      generateLogs(5),
+    );
+    const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
+    const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
+    const { result } = renderUseAgentLogs(queryClient, "ready");
+    await waitFor(() => {
+      expect(result.current).toHaveLength(5);
+    });
+    expect(fetchSpy).not.toHaveBeenCalled();
+    expect(wsSpy).not.toHaveBeenCalled();
+  });
+
+  it("should fetch logs when empty", async () => {
+    const queryClient = createTestQueryClient();
+    const fetchSpy = jest
+      .spyOn(API, "getWorkspaceAgentLogs")
+      .mockResolvedValueOnce(generateLogs(5));
+    const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
+    const { result } = renderUseAgentLogs(queryClient, "ready");
+    await waitFor(() => {
+      expect(result.current).toHaveLength(5);
+    });
+    expect(fetchSpy).toHaveBeenCalledWith(MockWorkspaceAgent.id);
+    expect(wsSpy).not.toHaveBeenCalled();
+  });
+
+  it("should fetch logs and connect to websocket when agent is starting", async () => {
+    const queryClient = createTestQueryClient();
+    const logs = generateLogs(5);
+    const fetchSpy = jest
+      .spyOn(API, "getWorkspaceAgentLogs")
+      .mockResolvedValueOnce(logs);
+    const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
+    new WS(
+      `ws://localhost/api/v2/workspaceagents/${
+        MockWorkspaceAgent.id
+      }/logs?follow&after=${logs[logs.length - 1].id}`,
+    );
+    const { result } = renderUseAgentLogs(queryClient, "starting");
+    await waitFor(() => {
+      expect(result.current).toHaveLength(5);
+    });
+    expect(fetchSpy).toHaveBeenCalledWith(MockWorkspaceAgent.id);
+    expect(wsSpy).toHaveBeenCalledWith(MockWorkspaceAgent.id, {
+      after: logs[logs.length - 1].id,
+      onMessage: expect.any(Function),
+      onError: expect.any(Function),
+    });
+  });
+
+  it("update logs from websocket messages", async () => {
+    const queryClient = createTestQueryClient();
+    const logs = generateLogs(5);
+    jest.spyOn(API, "getWorkspaceAgentLogs").mockResolvedValueOnce(logs);
+    const server = new WS(
+      `ws://localhost/api/v2/workspaceagents/${
+        MockWorkspaceAgent.id
+      }/logs?follow&after=${logs[logs.length - 1].id}`,
+    );
+    const { result } = renderUseAgentLogs(queryClient, "starting");
+    await waitFor(() => {
+      expect(result.current).toHaveLength(5);
+    });
+    await server.connected;
+    act(() => {
+      server.send(JSON.stringify(generateLogs(3)));
+    });
+    await waitFor(() => {
+      expect(result.current).toHaveLength(8);
+    });
+  });
+});
+
+function renderUseAgentLogs(
+  queryClient: QueryClient,
+  lifeCycleState: WorkspaceAgentLifecycle,
+  options?: UseAgentLogsOptions,
+) {
+  return renderHook(
+    () =>
+      useAgentLogs(
+        MockWorkspace.id,
+        MockWorkspaceAgent.id,
+        lifeCycleState,
+        options,
+      ),
+    {
+      wrapper: ({ children }) => (
+        <QueryClientProvider client={queryClient}>
+          {children}
+        </QueryClientProvider>
+      ),
+    },
+  );
+}
+
+function generateLogs(count: number): WorkspaceAgentLog[] {
+  return Array.from({ length: count }, (_, i) => ({
+    id: i,
+    created_at: new Date().toISOString(),
+    level: "info",
+    output: `Log ${i}`,
+    source_id: "",
+  }));
+}
diff --git a/site/src/modules/resources/AgentLogs/useAgentLogs.ts b/site/src/modules/resources/AgentLogs/useAgentLogs.ts
@@ -8,11 +8,15 @@ import type {
 } from "api/typesGenerated";
 import { useEffectEvent } from "hooks/hookPolyfills";
 
+export type UseAgentLogsOptions = {
+  enabled?: boolean;
+};
+
 export const useAgentLogs = (
   workspaceId: string,
   agentId: string,
   agentLifeCycleState: WorkspaceAgentLifecycle,
-  options?: { enabled?: boolean },
+  options?: UseAgentLogsOptions,
 ) => {
   const queryClient = useQueryClient();
   const queryOptions = agentLogs(workspaceId, agentId);