test: simplify mocks

app/server/modules/agents/__tests__/agents-manager.backups.test.ts

@@ -1,22 +1,13 @@
 import { afterEach, expect, test, vi } from "vitest";
 import waitForExpect from "wait-for-expect";
-import { fromPartial } from "@total-typescript/shoehorn";
-import { agentManager } from "../agents-manager";
+import { fromAny, fromPartial } from "@total-typescript/shoehorn";
+import { agentManager, type ProcessWithAgentRuntime } from "../agents-manager";
 import type { AgentManagerRuntime } from "../controller/server";
 import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
 
-type ProcessWithAgentRuntime = NodeJS.Process & {
-	__zerobyteAgentRuntime?: {
-		agentManager: AgentManagerRuntime | null;
-		localAgent: null;
-		isStoppingLocalAgent: boolean;
-		localAgentRestartTimeout: null;
-	};
-};
-
 const setAgentRuntime = (agentManagerRuntime: Partial<AgentManagerRuntime> | null) => {
 	(process as ProcessWithAgentRuntime).__zerobyteAgentRuntime = {
-		agentManager: agentManagerRuntime ? fromPartial<AgentManagerRuntime>(agentManagerRuntime) : null,
+		agentManager: fromAny(agentManagerRuntime),
 		localAgent: null,
 		isStoppingLocalAgent: false,
 		localAgentRestartTimeout: null,

app/server/modules/agents/__tests__/agents-manager.test.ts

@@ -1,16 +1,11 @@
 import { EventEmitter } from "node:events";
-import type * as childProcess from "node:child_process";
 import { PassThrough } from "node:stream";
 import { afterEach, expect, test, vi } from "vitest";
 
 const spawnMock = vi.fn();
 
 vi.mock("node:child_process", async () => {
-	const actual = await vi.importActual<typeof import("node:child_process")>("node:child_process");
-	return {
-		...actual,
-		spawn: spawnMock,
-	};
+	return { spawn: spawnMock };
 });
 
 const { spawnLocalAgent, stopLocalAgent } = await import("../agents-manager");
@@ -51,9 +46,7 @@ test("respawns the local agent after an unexpected exit", async () => {
 
 	const firstChild = createFakeChild();
 	const secondChild = createFakeChild();
-	spawnMock
-		.mockReturnValueOnce(firstChild as unknown as ReturnType<typeof childProcess.spawn>)
-		.mockReturnValueOnce(secondChild as unknown as ReturnType<typeof childProcess.spawn>);
+	spawnMock.mockReturnValueOnce(firstChild).mockReturnValueOnce(secondChild);
 
 	await spawnLocalAgent();
 
@@ -69,7 +62,7 @@ test("does not respawn the local agent after an intentional stop", async () => {
 	vi.useFakeTimers();
 
 	const child = createFakeChild();
-	spawnMock.mockReturnValue(child as unknown as ReturnType<typeof childProcess.spawn>);
+	spawnMock.mockReturnValue(child);
 
 	await spawnLocalAgent();
 	await stopLocalAgent();

app/server/modules/agents/__tests__/session.test.ts

@@ -6,19 +6,19 @@ import { createAgentMessage } from "@zerobyte/contracts/agent-protocol";
 import { createControllerAgentSession } from "../controller/session";
 
 const createSocket = (overrides: Partial<Parameters<typeof createControllerAgentSession>[0]> = {}) => {
-	return fromPartial<Parameters<typeof createControllerAgentSession>[0]>({
+	return {
 		data: { id: "connection-1", agentId: "local", organizationId: null, agentName: "Local Agent" },
 		send: vi.fn(() => 1),
 		close: vi.fn(),
 		...overrides,
-	});
+	};
 };
 
 const createSession = (handlers: Parameters<typeof createControllerAgentSession>[1] = {}, socket = createSocket()) => {
 	const scope = Effect.runSync(Scope.make());
 
 	try {
-		const session = Effect.runSync(Scope.extend(createControllerAgentSession(socket, handlers), scope));
+		const session = Effect.runSync(Scope.extend(createControllerAgentSession(fromPartial(socket), handlers), scope));
 
 		return {
 			session,
@@ -57,12 +57,12 @@ test("close emits a synthetic backup.cancelled for a started backup", () => {
 	close();
 
 	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
-	expect(onBackupCancelled).toHaveBeenCalledWith({
-		jobId: "job-1",
-		scheduleId: "schedule-1",
-		message:
-			"The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
-	});
+	expect(onBackupCancelled).toHaveBeenCalledWith(
+		expect.objectContaining({
+			jobId: "job-1",
+			scheduleId: "schedule-1",
+		}),
+	);
 });
 
 test.each([
@@ -150,12 +150,12 @@ test("close emits a synthetic backup.cancelled for a queued backup", () => {
 	close();
 
 	expect(onBackupCancelled).toHaveBeenCalledTimes(1);
-	expect(onBackupCancelled).toHaveBeenCalledWith({
-		jobId: "job-queued",
-		scheduleId: "schedule-queued",
-		message:
-			"The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.",
-	});
+	expect(onBackupCancelled).toHaveBeenLastCalledWith(
+		expect.objectContaining({
+			jobId: "job-queued",
+			scheduleId: "schedule-queued",
+		}),
+	);
 });
 
 test("a dropped backup.cancel closes the session and emits a synthetic backup.cancelled", async () => {
@@ -185,12 +185,12 @@ test("a dropped backup.cancel closes the session and emits a synthetic backup.ca
 			expect(send).toHaveBeenCalledTimes(1);
 			expect(socket.close).toHaveBeenCalledTimes(1);
 			expect(onBackupCancelled).toHaveBeenCalledTimes(1);
-			expect(onBackupCancelled).toHaveBeenCalledWith({
-				jobId: "job-1",
-				scheduleId: "schedule-1",
-				message:
-					"The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.",
-			});
+			expect(onBackupCancelled).toHaveBeenCalledWith(
+				expect.objectContaining({
+					jobId: "job-1",
+					scheduleId: "schedule-1",
+				}),
+			);
 		});
 	} finally {
 		await closeAsync();

app/server/modules/agents/agents-manager.ts

@@ -19,38 +19,6 @@ export type AgentRunBackupRequest = {
 	onProgress: (progress: BackupExecutionProgress) => void;
 };
 
-type AgentRuntimeState = {
-	agentManager: AgentManagerRuntime | null;
-	localAgent: ChildProcess | null;
-	isStoppingLocalAgent: boolean;
-	localAgentRestartTimeout: ReturnType<typeof setTimeout> | null;
-};
-
-type ProcessWithAgentRuntime = NodeJS.Process & {
-	__zerobyteAgentRuntime?: AgentRuntimeState;
-};
-
-const getAgentRuntimeState = () => {
-	const runtimeProcess = process as ProcessWithAgentRuntime;
-	const existingRuntime = runtimeProcess.__zerobyteAgentRuntime;
-
-	if (existingRuntime) {
-		return existingRuntime;
-	}
-
-	const runtime = {
-		agentManager: null,
-		localAgent: null,
-		isStoppingLocalAgent: false,
-		localAgentRestartTimeout: null,
-	};
-
-	runtimeProcess.__zerobyteAgentRuntime = runtime;
-	return runtime;
-};
-
-const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager;
-
 type ActiveBackupRun = {
 	scheduleId: number;
 	jobId: string;
@@ -60,8 +28,53 @@ type ActiveBackupRun = {
 	cancellationRequested: boolean;
 };
 
-const activeBackupsByScheduleId = new Map<number, ActiveBackupRun>();
-const activeBackupScheduleIdsByJobId = new Map<string, number>();
+type AgentRuntimeState = {
+	agentManager: AgentManagerRuntime | null;
+	localAgent: ChildProcess | null;
+	isStoppingLocalAgent: boolean;
+	localAgentRestartTimeout: ReturnType<typeof setTimeout> | null;
+	activeBackupsByScheduleId: Map<number, ActiveBackupRun>;
+	activeBackupScheduleIdsByJobId: Map<string, number>;
+};
+
+type LegacyAgentRuntimeState = Omit<AgentRuntimeState, "activeBackupsByScheduleId" | "activeBackupScheduleIdsByJobId"> &
+	Partial<Pick<AgentRuntimeState, "activeBackupsByScheduleId" | "activeBackupScheduleIdsByJobId">>;
+
+export type ProcessWithAgentRuntime = NodeJS.Process & {
+	__zerobyteAgentRuntime?: LegacyAgentRuntimeState;
+};
+
+const getAgentRuntimeState = () => {
+	const runtimeProcess = process as ProcessWithAgentRuntime;
+	const existingRuntime = runtimeProcess.__zerobyteAgentRuntime;
+
+	if (existingRuntime) {
+		const runtime: AgentRuntimeState = {
+			...existingRuntime,
+			activeBackupsByScheduleId: existingRuntime.activeBackupsByScheduleId ?? new Map<number, ActiveBackupRun>(),
+			activeBackupScheduleIdsByJobId: existingRuntime.activeBackupScheduleIdsByJobId ?? new Map<string, number>(),
+		};
+
+		runtimeProcess.__zerobyteAgentRuntime = runtime;
+		return runtime;
+	}
+
+	const runtime: AgentRuntimeState = {
+		agentManager: null,
+		localAgent: null,
+		isStoppingLocalAgent: false,
+		localAgentRestartTimeout: null,
+		activeBackupsByScheduleId: new Map(),
+		activeBackupScheduleIdsByJobId: new Map(),
+	};
+
+	runtimeProcess.__zerobyteAgentRuntime = runtime;
+	return runtime;
+};
+
+const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager;
+const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId;
+const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId;
 
 const getUnavailableError = (agentId: string) => {
 	if (agentId === "local") {
@@ -72,6 +85,8 @@ const getUnavailableError = (agentId: string) => {
 };
 
 const clearActiveBackupRun = (scheduleId: number) => {
+	const activeBackupsByScheduleId = getActiveBackupsByScheduleId();
+	const activeBackupScheduleIdsByJobId = getActiveBackupScheduleIdsByJobId();
 	const activeBackupRun = activeBackupsByScheduleId.get(scheduleId);
 	if (!activeBackupRun) {
 		return null;
@@ -93,13 +108,13 @@ const resolveActiveBackupRun = (scheduleId: number, result: BackupExecutionResul
 };
 
 const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string, agentId: string) => {
-	const trackedScheduleId = activeBackupScheduleIdsByJobId.get(jobId);
+	const trackedScheduleId = getActiveBackupScheduleIdsByJobId().get(jobId);
 	if (trackedScheduleId === undefined) {
 		logger.warn(`Received ${eventName} for unknown job ${jobId} from agent ${agentId}`);
 		return null;
 	}
 
-	const activeBackupRun = activeBackupsByScheduleId.get(trackedScheduleId);
+	const activeBackupRun = getActiveBackupsByScheduleId().get(trackedScheduleId);
 	if (!activeBackupRun) {
 		logger.warn(`Received ${eventName} for inactive job ${jobId} from agent ${agentId}`);
 		return null;
@@ -114,7 +129,7 @@ const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string
 };
 
 const requestBackupCancellation = async (agentId: string, scheduleId: number) => {
-	const activeBackupRun = activeBackupsByScheduleId.get(scheduleId);
+	const activeBackupRun = getActiveBackupsByScheduleId().get(scheduleId);
 	if (!activeBackupRun) {
 		return false;
 	}
@@ -223,7 +238,7 @@ export const agentManager = {
 		}
 
 		const completion = new Promise<BackupExecutionResult>((resolve) => {
-			activeBackupsByScheduleId.set(request.scheduleId, {
+			getActiveBackupsByScheduleId().set(request.scheduleId, {
 				scheduleId: request.scheduleId,
 				jobId: request.payload.jobId,
 				scheduleShortId: request.payload.scheduleId,
@@ -231,7 +246,7 @@ export const agentManager = {
 				resolve,
 				cancellationRequested: false,
 			});
-			activeBackupScheduleIdsByJobId.set(request.payload.jobId, request.scheduleId);
+			getActiveBackupScheduleIdsByJobId().set(request.payload.jobId, request.scheduleId);
 		});
 
 		try {