diff --git a/site/src/modules/resources/useAgentLogs.test.ts b/site/src/modules/resources/useAgentLogs.test.ts index a5339e00c87eb..85206171bf7d5 100644 --- a/site/src/modules/resources/useAgentLogs.test.ts +++ b/site/src/modules/resources/useAgentLogs.test.ts @@ -1,60 +1,114 @@ import { renderHook, waitFor } from "@testing-library/react"; import type { WorkspaceAgentLog } from "api/typesGenerated"; -import WS from "jest-websocket-mock"; import { MockWorkspaceAgent } from "testHelpers/entities"; -import { useAgentLogs } from "./useAgentLogs"; +import { + type MockWebSocketPublisher, + createMockWebSocket, +} from "testHelpers/websockets"; +import { OneWayWebSocket } from "utils/OneWayWebSocket"; +import { createUseAgentLogs } from "./useAgentLogs"; -/** - * TODO: WS does not support multiple tests running at once in isolation so we - * have one single test that test the most common scenario. - * Issue: https://github.com/romgain/jest-websocket-mock/issues/172 - */ +const millisecondsInOneMinute = 60_000; -describe("useAgentLogs", () => { - afterEach(() => { - WS.clean(); +function generateMockLogs( + logCount: number, + baseDate = new Date(), +): readonly WorkspaceAgentLog[] { + return Array.from({ length: logCount }, (_, i) => { + // Make sure that the logs generated each have unique timestamps, so + // that we can test whether they're being sorted properly before being + // returned by the hook + const logDate = new Date(baseDate.getTime() + i * millisecondsInOneMinute); + return { + id: i, + created_at: logDate.toISOString(), + level: "info", + output: `Log ${i}`, + source_id: "", + }; }); +} - it("clear logs when disabled to avoid duplicates", async () => { - const server = new WS( - `ws://localhost/api/v2/workspaceagents/${ - MockWorkspaceAgent.id - }/logs?follow&after=0`, - ); - const { result, rerender } = renderHook( - ({ enabled }) => useAgentLogs(MockWorkspaceAgent, enabled), - { initialProps: { enabled: true } }, - ); - await server.connected; - - // Send 3 logs - server.send(JSON.stringify(generateLogs(3))); - await waitFor(() => { - expect(result.current).toHaveLength(3); +// A mutable object holding the most recent mock WebSocket publisher. The inner +// value will change as the hook opens/closes new connections +type PublisherResult = { + current: MockWebSocketPublisher; +}; + +type MountHookResult = Readonly<{ + // Note: the value of `current` should be readonly, but the `current` + // property itself should be mutable + hookResult: { + current: readonly WorkspaceAgentLog[]; + }; + rerender: (props: { enabled: boolean }) => void; + publisherResult: PublisherResult; +}>; + +function mountHook(): MountHookResult { + // Have to cheat the types a little bit to avoid a chicken-and-the-egg + // scenario. publisherResult will be initialized with an undefined current + // value, but it'll be guaranteed not to be undefined by the time this + // function returns. + const publisherResult: Partial = { current: undefined }; + const useAgentLogs = createUseAgentLogs((agentId, params) => { + return new OneWayWebSocket({ + apiRoute: `/api/v2/workspaceagents/${agentId}/logs`, + searchParams: new URLSearchParams({ + follow: "true", + after: params?.after?.toString() || "0", + }), + websocketInit: (url) => { + const [mockSocket, mockPublisher] = createMockWebSocket(url); + publisherResult.current = mockPublisher; + return mockSocket; + }, }); + }); - // Disable the hook - rerender({ enabled: false }); + const { result, rerender } = renderHook( + ({ enabled }) => useAgentLogs(MockWorkspaceAgent, enabled), + { initialProps: { enabled: true } }, + ); + + return { + rerender, + hookResult: result, + publisherResult: publisherResult as PublisherResult, + }; +} + +describe("useAgentLogs", () => { + it("clears logs when hook becomes disabled (protection to avoid duplicate logs when hook goes back to being re-enabled)", async () => { + const { hookResult, publisherResult, rerender } = mountHook(); + + // Verify that logs can be received after mount + const initialLogs = generateMockLogs(3, new Date("april 5, 1997")); + const initialEvent = new MessageEvent("message", { + data: JSON.stringify(initialLogs), + }); + publisherResult.current.publishMessage(initialEvent); await waitFor(() => { - expect(result.current).toHaveLength(0); + // Using expect.arrayContaining to account for the fact that we're + // not guaranteed to receive WebSocket events in order + expect(hookResult.current).toEqual(expect.arrayContaining(initialLogs)); }); - // Enable the hook again + // Disable the hook (and have the hook close the connection behind the + // scenes) + rerender({ enabled: false }); + await waitFor(() => expect(hookResult.current).toHaveLength(0)); + + // Re-enable the hook (creating an entirely new connection), and send + // new logs rerender({ enabled: true }); - await server.connected; - server.send(JSON.stringify(generateLogs(3))); + const newLogs = generateMockLogs(3, new Date("october 3, 2005")); + const newEvent = new MessageEvent("message", { + data: JSON.stringify(newLogs), + }); + publisherResult.current.publishMessage(newEvent); await waitFor(() => { - expect(result.current).toHaveLength(3); + expect(hookResult.current).toEqual(expect.arrayContaining(newLogs)); }); }); }); - -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/useAgentLogs.ts b/site/src/modules/resources/useAgentLogs.ts index d7f810483a693..63ea4afbb99c1 100644 --- a/site/src/modules/resources/useAgentLogs.ts +++ b/site/src/modules/resources/useAgentLogs.ts @@ -3,45 +3,63 @@ import type { WorkspaceAgent, WorkspaceAgentLog } from "api/typesGenerated"; import { displayError } from "components/GlobalSnackbar/utils"; import { useEffect, useState } from "react"; -export function useAgentLogs( - agent: WorkspaceAgent, - enabled: boolean, -): readonly WorkspaceAgentLog[] { - const [logs, setLogs] = useState([]); - - useEffect(() => { - if (!enabled) { - // Clean up the logs when the agent is not enabled. So it can receive logs - // from the beginning without duplicating the logs. +export function createUseAgentLogs( + createSocket: typeof watchWorkspaceAgentLogs, +) { + return function useAgentLogs( + agent: WorkspaceAgent, + enabled: boolean, + ): readonly WorkspaceAgentLog[] { + const [logs, setLogs] = useState([]); + + // Clean up the logs when the agent is not enabled, using a mid-render + // sync to remove any risk of screen flickering. Clearing the logs helps + // ensure that if the hook flips back to being enabled, we can receive a + // fresh set of logs from the beginning with zero risk of duplicates. + const [prevEnabled, setPrevEnabled] = useState(enabled); + if (!enabled && prevEnabled) { setLogs([]); - return; + setPrevEnabled(false); } - // Always fetch the logs from the beginning. We may want to optimize this in - // the future, but it would add some complexity in the code that maybe does - // not worth it. - const socket = watchWorkspaceAgentLogs(agent.id, { after: 0 }); - socket.addEventListener("message", (e) => { - if (e.parseError) { - console.warn("Error parsing agent log: ", e.parseError); + useEffect(() => { + if (!enabled) { return; } - setLogs((logs) => [...logs, ...e.parsedMessage]); - }); - - socket.addEventListener("error", (e) => { - console.error("Error in agent log socket: ", e); - displayError( - "Unable to watch the agent logs", - "Please try refreshing the browser", - ); - socket.close(); - }); - - return () => { - socket.close(); - }; - }, [agent.id, enabled]); - - return logs; + + // Always fetch the logs from the beginning. We may want to optimize + // this in the future, but it would add some complexity in the code + // that might not be worth it. + const socket = createSocket(agent.id, { after: 0 }); + socket.addEventListener("message", (e) => { + if (e.parseError) { + console.warn("Error parsing agent log: ", e.parseError); + return; + } + setLogs((logs) => [...logs, ...e.parsedMessage]); + }); + + socket.addEventListener("error", (e) => { + console.error("Error in agent log socket: ", e); + displayError( + "Unable to watch the agent logs", + "Please try refreshing the browser", + ); + socket.close(); + }); + + return () => socket.close(); + + // createSocket shouldn't ever change for the lifecycle of the hook, + // but Biome isn't smart enough to detect constant dependencies for + // higher-order hooks. Adding it to the array (even though it + // shouldn't ever be needed) seemed like the least fragile way to + // resolve the warning. + }, [createSocket, agent.id, enabled]); + + return logs; + }; } + +// The baseline implementation to use for production +export const useAgentLogs = createUseAgentLogs(watchWorkspaceAgentLogs); diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts new file mode 100644 index 0000000000000..ac356f1c8dc28 --- /dev/null +++ b/site/src/testHelpers/websockets.ts @@ -0,0 +1,135 @@ +import type { WebSocketEventType } from "utils/OneWayWebSocket"; + +export type MockWebSocketPublisher = Readonly<{ + publishMessage: (event: MessageEvent) => void; + publishError: (event: ErrorEvent) => void; + publishClose: (event: CloseEvent) => void; + publishOpen: (event: Event) => void; +}>; + +export function createMockWebSocket( + url: string, + protocols?: string | string[], +): readonly [WebSocket, MockWebSocketPublisher] { + type EventMap = { + message: MessageEvent; + error: ErrorEvent; + close: CloseEvent; + open: Event; + }; + type CallbackStore = { + [K in keyof EventMap]: ((event: EventMap[K]) => void)[]; + }; + + let activeProtocol: string; + if (Array.isArray(protocols)) { + activeProtocol = protocols[0] ?? ""; + } else if (typeof protocols === "string") { + activeProtocol = protocols; + } else { + activeProtocol = ""; + } + + let closed = false; + const store: CallbackStore = { + message: [], + error: [], + close: [], + open: [], + }; + + const mockSocket: WebSocket = { + CONNECTING: 0, + OPEN: 1, + CLOSING: 2, + CLOSED: 3, + + url, + protocol: activeProtocol, + readyState: 1, + binaryType: "blob", + bufferedAmount: 0, + extensions: "", + onclose: null, + onerror: null, + onmessage: null, + onopen: null, + send: jest.fn(), + dispatchEvent: jest.fn(), + + addEventListener: ( + eventType: E, + callback: WebSocketEventMap[E], + ) => { + if (closed) { + return; + } + + const subscribers = store[eventType]; + const cb = callback as unknown as CallbackStore[E][0]; + if (!subscribers.includes(cb)) { + subscribers.push(cb); + } + }, + + removeEventListener: ( + eventType: E, + callback: WebSocketEventMap[E], + ) => { + if (closed) { + return; + } + + const subscribers = store[eventType]; + const cb = callback as unknown as CallbackStore[E][0]; + if (subscribers.includes(cb)) { + const updated = store[eventType].filter((c) => c !== cb); + store[eventType] = updated as unknown as CallbackStore[E]; + } + }, + + close: () => { + closed = true; + }, + }; + + const publisher: MockWebSocketPublisher = { + publishOpen: (event) => { + if (closed) { + return; + } + for (const sub of store.open) { + sub(event); + } + }, + + publishError: (event) => { + if (closed) { + return; + } + for (const sub of store.error) { + sub(event); + } + }, + + publishMessage: (event) => { + if (closed) { + return; + } + for (const sub of store.message) { + sub(event); + } + }, + + publishClose: (event) => { + if (closed) { + return; + } + for (const sub of store.close) { + sub(event); + } + }, + }; + + return [mockSocket, publisher] as const; +} diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts index c6b00b593111f..b334732c5b5e8 100644 --- a/site/src/utils/OneWayWebSocket.test.ts +++ b/site/src/utils/OneWayWebSocket.test.ts @@ -8,144 +8,10 @@ */ import { - type OneWayMessageEvent, - OneWayWebSocket, - type WebSocketEventType, -} from "./OneWayWebSocket"; - -type MockPublisher = Readonly<{ - publishMessage: (event: MessageEvent) => void; - publishError: (event: ErrorEvent) => void; - publishClose: (event: CloseEvent) => void; - publishOpen: (event: Event) => void; -}>; - -function createMockWebSocket( - url: string, - protocols?: string | string[], -): readonly [WebSocket, MockPublisher] { - type EventMap = { - message: MessageEvent; - error: ErrorEvent; - close: CloseEvent; - open: Event; - }; - type CallbackStore = { - [K in keyof EventMap]: ((event: EventMap[K]) => void)[]; - }; - - let activeProtocol: string; - if (Array.isArray(protocols)) { - activeProtocol = protocols[0] ?? ""; - } else if (typeof protocols === "string") { - activeProtocol = protocols; - } else { - activeProtocol = ""; - } - - let closed = false; - const store: CallbackStore = { - message: [], - error: [], - close: [], - open: [], - }; - - const mockSocket: WebSocket = { - CONNECTING: 0, - OPEN: 1, - CLOSING: 2, - CLOSED: 3, - - url, - protocol: activeProtocol, - readyState: 1, - binaryType: "blob", - bufferedAmount: 0, - extensions: "", - onclose: null, - onerror: null, - onmessage: null, - onopen: null, - send: jest.fn(), - dispatchEvent: jest.fn(), - - addEventListener: ( - eventType: E, - callback: WebSocketEventMap[E], - ) => { - if (closed) { - return; - } - - const subscribers = store[eventType]; - const cb = callback as unknown as CallbackStore[E][0]; - if (!subscribers.includes(cb)) { - subscribers.push(cb); - } - }, - - removeEventListener: ( - eventType: E, - callback: WebSocketEventMap[E], - ) => { - if (closed) { - return; - } - - const subscribers = store[eventType]; - const cb = callback as unknown as CallbackStore[E][0]; - if (subscribers.includes(cb)) { - const updated = store[eventType].filter((c) => c !== cb); - store[eventType] = updated as unknown as CallbackStore[E]; - } - }, - - close: () => { - closed = true; - }, - }; - - const publisher: MockPublisher = { - publishOpen: (event) => { - if (closed) { - return; - } - for (const sub of store.open) { - sub(event); - } - }, - - publishError: (event) => { - if (closed) { - return; - } - for (const sub of store.error) { - sub(event); - } - }, - - publishMessage: (event) => { - if (closed) { - return; - } - for (const sub of store.message) { - sub(event); - } - }, - - publishClose: (event) => { - if (closed) { - return; - } - for (const sub of store.close) { - sub(event); - } - }, - }; - - return [mockSocket, publisher] as const; -} + type MockWebSocketPublisher, + createMockWebSocket, +} from "testHelpers/websockets"; +import { type OneWayMessageEvent, OneWayWebSocket } from "./OneWayWebSocket"; describe(OneWayWebSocket.name, () => { const dummyRoute = "/api/v2/blah"; @@ -167,7 +33,7 @@ describe(OneWayWebSocket.name, () => { }); it("Lets a consumer add an event listener of each type", () => { - let publisher!: MockPublisher; + let publisher!: MockWebSocketPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { @@ -207,7 +73,7 @@ describe(OneWayWebSocket.name, () => { }); it("Lets a consumer remove an event listener of each type", () => { - let publisher!: MockPublisher; + let publisher!: MockWebSocketPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { @@ -252,7 +118,7 @@ describe(OneWayWebSocket.name, () => { }); it("Only calls each callback once if callback is added multiple times", () => { - let publisher!: MockPublisher; + let publisher!: MockWebSocketPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { @@ -294,7 +160,7 @@ describe(OneWayWebSocket.name, () => { }); it("Lets consumers register multiple callbacks for each event type", () => { - let publisher!: MockPublisher; + let publisher!: MockWebSocketPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { @@ -375,7 +241,7 @@ describe(OneWayWebSocket.name, () => { }); it("Gives consumers pre-parsed versions of message events", () => { - let publisher!: MockPublisher; + let publisher!: MockWebSocketPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => { @@ -405,7 +271,7 @@ describe(OneWayWebSocket.name, () => { }); it("Exposes parsing error if message payload could not be parsed as JSON", () => { - let publisher!: MockPublisher; + let publisher!: MockWebSocketPublisher; const oneWay = new OneWayWebSocket({ apiRoute: dummyRoute, websocketInit: (url, protocols) => {