zerobyte/app/client/hooks/__tests__/use-server-events.test.tsx

import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
import { useEffect, useState } from "react";
import { useQuery } from "@tanstack/react-query";
import { cleanup, createTestQueryClient, render, screen } from "~/test/test-utils";
import { useServerEvents } from "../use-server-events";

class MockEventSource {
	static instances: MockEventSource[] = [];

	public onerror: ((event: Event) => void) | null = null;
	public close = vi.fn(() => {});
	private listeners = new Map<string, Set<(event: Event) => void>>();

	constructor(public url: string) {
		MockEventSource.instances.push(this);
	}

	addEventListener(type: string, listener: EventListenerOrEventListenerObject) {
		const listeners = this.listeners.get(type) ?? new Set<(event: Event) => void>();
		const callback = typeof listener === "function" ? listener : (event: Event) => listener.handleEvent(event);
		listeners.add(callback);
		this.listeners.set(type, listeners);
	}

	emit(type: string, data: unknown) {
		const event = new MessageEvent(type, {
			data: JSON.stringify(data),
		});
		for (const listener of this.listeners.get(type) ?? []) {
			listener(event);
		}
	}

	static reset() {
		MockEventSource.instances = [];
	}
}

const originalEventSource = globalThis.EventSource;
const originalConsoleInfo = console.info;
const originalConsoleError = console.error;

const ConnectionConsumer = ({ enabled = true }: { enabled?: boolean }) => {
	useServerEvents({ enabled });
	return null;
};

const BackupCompletedListener = ({ scheduleId }: { scheduleId: string }) => {
	const { addEventListener } = useServerEvents();
	const [status, setStatus] = useState("pending");

	useEffect(() => {
		const abortController = new AbortController();

		addEventListener(
			"backup:completed",
			(event) => {
				if (event.scheduleId === scheduleId) {
					setStatus(event.status);
				}
			},
			{ signal: abortController.signal },
		);

		return () => abortController.abort();
	}, [addEventListener, scheduleId]);

	return <div>{status}</div>;
};

const QueryStatusConsumer = ({ getValue }: { getValue: () => string }) => {
	const { data } = useQuery({
		queryKey: ["backup-status"],
		queryFn: async () => getValue(),
	});

	return <div>{data}</div>;
};

describe("useServerEvents", () => {
	beforeEach(() => {
		MockEventSource.reset();
		globalThis.EventSource = MockEventSource as unknown as typeof EventSource;
		console.info = vi.fn(() => {});
		console.error = vi.fn(() => {});
	});

	afterEach(() => {
		cleanup();
		globalThis.EventSource = originalEventSource;
		console.info = originalConsoleInfo;
		console.error = originalConsoleError;
		MockEventSource.reset();
	});

	test("shares one EventSource across consumers and refreshes active queries on backup completion", async () => {
		const queryClient = createTestQueryClient();
		let queryValue = "before";

		const view = render(
			<>
				<ConnectionConsumer />
				<BackupCompletedListener scheduleId="0b9c940b" />
				<QueryStatusConsumer getValue={() => queryValue} />
			</>,
			{ queryClient },
		);

		expect(MockEventSource.instances).toHaveLength(1);
		expect(await screen.findByText("before")).toBeTruthy();

		queryValue = "after";

		MockEventSource.instances[0]?.emit("backup:completed", {
			organizationId: "default-org",
			scheduleId: "0b9c940b",
			volumeName: "synology",
			repositoryName: "swiss-backup",
			status: "success",
		});

		expect(await screen.findByText("success")).toBeTruthy();
		expect(await screen.findByText("after")).toBeTruthy();

		view.unmount();

		expect(MockEventSource.instances[0]?.close).toHaveBeenCalledTimes(1);
	});

	test("waits to subscribe until enabled", () => {
		const queryClient = createTestQueryClient();
		const view = render(<ConnectionConsumer enabled={false} />, { queryClient });

		expect(MockEventSource.instances).toHaveLength(0);

		view.rerender(<ConnectionConsumer />);

		expect(MockEventSource.instances).toHaveLength(1);
		expect(MockEventSource.instances[0]?.url).toBe("/api/v1/events");
	});
});