chore: add one more test case

Parkreiner · Parkreiner · commit 550d09e422a8 · 2025-08-02T20:38:56.000Z
diff --git a/site/src/modules/resources/useAgentLogs.test.ts b/site/src/modules/resources/useAgentLogs.test.ts
@@ -17,8 +17,8 @@ function generateMockLogs(
 ): 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
+		// that we can test whether the hook is sorting them properly as it's
+		// receiving them over time
 		const logDate = new Date(baseDate.getTime() + i * millisecondsInOneMinute);
 		return {
 			id: i,
@@ -32,25 +32,22 @@ function generateMockLogs(
 
 // 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 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;
+	rerender: (props: { agentId: string, enabled: boolean }) => void;
 	publisherResult: PublisherResult;
+
+	// Note: the `current` property is only "halfway" readonly; the value is
+	// readonly, but the key is still mutable
+	hookResult: { current: readonly WorkspaceAgentLog[]; };
 }>;
 
-function mountHook(): MountHookResult {
+function mountHook(initialAgentId: string): 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.
+	// value, but it'll be guaranteed not to be undefined by the time mountHook
+	// returns its value.
 	const publisherResult: Partial<PublisherResult> = { current: undefined };
 	const useAgentLogs = createUseAgentLogs((agentId, params) => {
 		return new OneWayWebSocket({
@@ -68,8 +65,8 @@ function mountHook(): MountHookResult {
 	});
 
 	const { result, rerender } = renderHook(
-		({ enabled }) => useAgentLogs(MockWorkspaceAgent, enabled),
-		{ initialProps: { enabled: true } },
+		({ agentId, enabled }) => useAgentLogs(agentId, enabled),
+		{ initialProps: { agentId: initialAgentId, enabled: true }, },
 	);
 
 	return {
@@ -81,10 +78,11 @@ function mountHook(): MountHookResult {
 
 describe("useAgentLogs", () => {
 	it("Automatically sorts logs that are received out of order", async () => {
-		const { hookResult, publisherResult } = mountHook();
+		const { hookResult, publisherResult } = mountHook(MockWorkspaceAgent.id);
 		const logs = generateMockLogs(10, new Date("september 9, 1999"));
+		const reversed = logs.toReversed();
 
-		for (const log of logs.toReversed()) {
+		for (const log of reversed) {
 			act(() => {
 				publisherResult.current.publishMessage(
 					new MessageEvent<string>("message", {
@@ -97,16 +95,30 @@ describe("useAgentLogs", () => {
 	});
 
 	it("Automatically closes the socket connection when the hook is disabled", async () => {
-		const { publisherResult, rerender } = mountHook();
+		const { publisherResult, rerender } = mountHook(MockWorkspaceAgent.id);
 		expect(publisherResult.current.isConnectionOpen()).toBe(true);
-		rerender({ enabled: false });
+		rerender({ agentId: MockWorkspaceAgent.id, enabled: false });
 		await waitFor(() => {
 			expect(publisherResult.current.isConnectionOpen()).toBe(false);
 		})
 	});
 
+	it("Automatically closes the old connection when the agent ID changes", () => {
+		const { publisherResult, rerender } = mountHook(MockWorkspaceAgent.id);
+		const publisher1 = publisherResult.current;
+		expect(publisher1.isConnectionOpen()).toBe(true);
+
+		const newAgentId = `${MockWorkspaceAgent.id}-2`;
+		rerender({ agentId: newAgentId, enabled: true });
+
+		const publisher2 = publisherResult.current;
+		expect(publisher1.isConnectionOpen()).toBe(false);
+		expect(publisher2.isConnectionOpen()).toBe(true);
+	});
+
 	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();
+		const { hookResult, publisherResult, rerender }
+			= mountHook(MockWorkspaceAgent.id);
 
 		// Send initial logs so that we have something to clear out later
 		const initialLogs = generateMockLogs(3, new Date("april 5, 1997"));
@@ -118,12 +130,12 @@ describe("useAgentLogs", () => {
 
 		// Disable the hook (and have the hook close the connection behind the
 		// scenes)
-		rerender({ enabled: false });
+		rerender({ agentId: MockWorkspaceAgent.id, 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 });
+		rerender({ agentId: MockWorkspaceAgent.id, enabled: true });
 		const newLogs = generateMockLogs(3, new Date("october 3, 2005"));
 		const newEvent = new MessageEvent<string>("message", {
 			data: JSON.stringify(newLogs),
diff --git a/site/src/modules/resources/useAgentLogs.ts b/site/src/modules/resources/useAgentLogs.ts
@@ -1,5 +1,5 @@
 import { watchWorkspaceAgentLogs, type WatchWorkspaceAgentLogsParams } from "api/api";
-import type { WorkspaceAgent, WorkspaceAgentLog } from "api/typesGenerated";
+import type { WorkspaceAgentLog } from "api/typesGenerated";
 import { displayError } from "components/GlobalSnackbar/utils";
 import { useEffect, useState } from "react";
 import type { OneWayWebSocket } from "utils/OneWayWebSocket";
@@ -11,7 +11,7 @@ export type CreateSocket = (
 
 export function createUseAgentLogs(createSocket: CreateSocket) {
 	return function useAgentLogs(
-		agent: WorkspaceAgent,
+		agentId: string,
 		enabled: boolean,
 	): readonly WorkspaceAgentLog[] {
 		const [logs, setLogs] = useState<readonly WorkspaceAgentLog[]>([]);
@@ -35,7 +35,7 @@ export function createUseAgentLogs(createSocket: CreateSocket) {
 			// this in the future, but it would add some complexity in the code
 			// that might not be worth it.
 			const createdAtMap = new Map<string, number>();
-			const socket = createSocket(agent.id, { after: 0 });
+			const socket = createSocket(agentId, { after: 0 });
 			socket.addEventListener("message", (e) => {
 				if (e.parseError) {
 					console.warn("Error parsing agent log: ", e.parseError);
@@ -76,7 +76,7 @@ export function createUseAgentLogs(createSocket: CreateSocket) {
 			});
 
 			return () => socket.close();
-		}, [createSocket, agent.id, enabled]);
+		}, [createSocket, agentId, enabled]);
 
 		return logs;
 	};
