refactor: update useAgentLogs tests as unit tests

Parkreiner · Parkreiner · commit 43d0ca8ec095 · 2025-05-23T20:25:25.000Z
diff --git a/site/src/modules/resources/useAgentLogs.test.ts b/site/src/modules/resources/useAgentLogs.test.ts
@@ -1,8 +1,4 @@
-import {
-	renderHook,
-	type RenderHookResult,
-	waitFor,
-} from "@testing-library/react";
+import { renderHook, waitFor } from "@testing-library/react";
 import type { WorkspaceAgentLog } from "api/typesGenerated";
 import { MockWorkspaceAgent } from "testHelpers/entities";
 import { createUseAgentLogs } from "./useAgentLogs";
@@ -22,14 +18,28 @@ function generateMockLogs(count: number): WorkspaceAgentLog[] {
 	}));
 }
 
-type MountHookResult = Readonly<
-	RenderHookResult<readonly WorkspaceAgentLog[], { enabled: boolean }> & {
-		publisher: MockWebSocketPublisher;
-	}
->;
+// A mutable object holding the most recent mock WebSocket connection. This
+// 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 {
-	let publisher!: MockWebSocketPublisher;
+	// 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<PublisherResult> = { current: undefined };
 	const useAgentLogs = createUseAgentLogs((agentId, params) => {
 		return new OneWayWebSocket({
 			apiRoute: `/api/v2/workspaceagents/${agentId}/logs`,
@@ -38,41 +48,45 @@ function mountHook(): MountHookResult {
 				after: params?.after?.toString() || "0",
 			}),
 			websocketInit: (url) => {
-				const [mockSocket, mockPub] = createMockWebSocket(url);
-				publisher = mockPub;
+				const [mockSocket, mockPublisher] = createMockWebSocket(url);
+				publisherResult.current = mockPublisher;
 				return mockSocket;
 			},
 		});
 	});
 
-	const { result, rerender, unmount } = renderHook(
+	const { result, rerender } = renderHook(
 		({ enabled }) => useAgentLogs(MockWorkspaceAgent, enabled),
 		{ initialProps: { enabled: true } },
 	);
 
-	return { result, rerender, unmount, publisher };
+	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 { result, rerender, publisher } = mountHook();
+		const { hookResult, publisherResult, rerender } = mountHook();
 
 		// Verify that logs can be received after mount
 		const initialLogs = generateMockLogs(3);
 		const initialEvent = new MessageEvent<string>("message", {
 			data: JSON.stringify(initialLogs),
 		});
-		publisher.publishMessage(initialEvent);
+		publisherResult.current.publishMessage(initialEvent);
 		await waitFor(() => {
 			// Using expect.arrayContaining to account for the fact that we're
 			// not guaranteed to receive WebSocket events in order
-			expect(result.current).toEqual(expect.arrayContaining(initialLogs));
+			expect(hookResult.current).toEqual(expect.arrayContaining(initialLogs));
 		});
 
 		// Disable the hook (and have the hook close the connection behind the
 		// scenes)
 		rerender({ enabled: false });
-		await waitFor(() => expect(result.current).toHaveLength(0));
+		await waitFor(() => expect(hookResult.current).toHaveLength(0));
 
 		// Re-enable the hook (creating an entirely new connection), and send
 		// new logs
@@ -81,9 +95,9 @@ describe("useAgentLogs", () => {
 		const newEvent = new MessageEvent<string>("message", {
 			data: JSON.stringify(newLogs),
 		});
-		publisher.publishMessage(newEvent);
+		publisherResult.current.publishMessage(newEvent);
 		await waitFor(() => {
-			expect(result.current).toEqual(expect.arrayContaining(newLogs));
+			expect(hookResult.current).toEqual(expect.arrayContaining(newLogs));
 		});
 	});
 });
