KSP-MissionControl/apps/api/tests/ws.test.ts

import { describe, it, expect } from 'vitest';
import { createServer, type Server } from 'node:http';
import { WebSocket } from 'ws';
import { createNodeWebSocket } from '@hono/node-ws';
import { InMemoryStateStore } from '@kerbal-rt/db';
import { buildApp } from '../src/app.js';
import { LiveSocketHub } from '../src/ws.js';

/**
 * Integration test: the API exposes /api/v1/live over WebSocket.
 * When a snapshot is applied to the store, the WS hub broadcasts it
 * to every connected client.
 */
describe('WebSocket live feed', () => {
  async function startTestServer(): Promise<{
    server: Server;
    port: number;
    store: InMemoryStateStore;
    close: () => Promise<void>;
  }> {
    const store = new InMemoryStateStore();
    const hub = new LiveSocketHub({ store, channel: 'test' });
    await hub.start();

    const app = buildApp({ store, ingestApiKey: null });
    const { upgradeWebSocket, injectWebSocket } = createNodeWebSocket({ app });

    app.get(
      '/api/v1/live',
      upgradeWebSocket(() => ({
        onOpen: (_evt, ws) => hub.onOpen(ws),
        onMessage: (evt, ws) => hub.onMessage(evt, ws),
        onClose: (_evt, ws) => hub.onClose(ws),
      })),
    );

    const server = createServer();
    injectWebSocket(server);
    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
    const port = (server.address() as { port: number }).port;

    return {
      server,
      port,
      store,
      async close() {
        await hub.stop();
        await new Promise<void>((resolve) => server.close(() => resolve()));
      },
    };
  }

  it('sends the current snapshot to a newly connected client', async () => {
    const { port, store, close } = await startTestServer();
    try {
      await store.applySnapshot({
        ut: 42,
        capturedAt: 'x',
        activeVesselId: null,
        bodies: [],
        vessels: [],
        groundStations: [],
      });

      const ws = new WebSocket(`ws://127.0.0.1:${port}/api/v1/live`);
      const firstMessage = await new Promise<string>((resolve, reject) => {
        ws.once('open', () => {});
        ws.once('message', (data) => {
          resolve(data.toString());
          ws.close();
        });
        ws.once('error', reject);
        setTimeout(() => reject(new Error('timeout')), 5000);
      });

      const parsed = JSON.parse(firstMessage);
      expect(parsed.type).toBe('snapshot');
      expect(parsed.snapshot.ut).toBe(42);
    } finally {
      await close();
    }
  });

  it('broadcasts new snapshots to all connected clients', async () => {
    const { port, store, close } = await startTestServer();
    try {
      const ws1 = new WebSocket(`ws://127.0.0.1:${port}/api/v1/live`);
      const ws2 = new WebSocket(`ws://127.0.0.1:${port}/api/v1/live`);

      // Wait for both to open
      await Promise.all([
        new Promise<void>((r) => ws1.once('open', () => r())),
        new Promise<void>((r) => ws2.once('open', () => r())),
      ]);

      // Attach persistent handlers FIRST (so we don't miss any message
      // during the drain→wait window), then drain the initial snapshot
      // messages and wait for the broadcast.
      const received = { ws1: false, ws2: false };
      const waiters: Promise<void>[] = [];
      const makeWaiter = (key: 'ws1' | 'ws2', ws: WebSocket) =>
        new Promise<void>((resolve) => {
          ws.on('message', (data) => {
            const msg = JSON.parse(data.toString());
            if (msg.type === 'snapshot' && msg.snapshot.ut === 99) {
              received[key] = true;
              resolve();
            }
          });
          // Safety timeout so a failing test doesn't hang forever
          setTimeout(() => resolve(), 3000);
        });
      waiters.push(makeWaiter('ws1', ws1));
      waiters.push(makeWaiter('ws2', ws2));

      // Wait for the initial snapshot to land on both before we apply
      // the new one — otherwise we may set up the waiter after it
      // arrives (race condition with `once`).
      await new Promise((r) => setTimeout(r, 100));

      await store.applySnapshot({
        ut: 99,
        capturedAt: 'x',
        activeVesselId: null,
        bodies: [],
        vessels: [],
        groundStations: [],
      });

      await Promise.all(waiters);

      expect(received.ws1).toBe(true);
      expect(received.ws2).toBe(true);

      ws1.close();
      ws2.close();
    } finally {
      await close();
    }
  });
});