lab/bastion/src/shared/tests/protocol.test.ts

// Tests for WebSocket protocol types, type guards, and parsing.

import { describe, it, expect } from "vitest";
import {
  isAgentMessage,
  isServerMessage,
  parseAgentMessage,
  parseServerMessage,
  generateRequestId,
} from "../src/protocol/index.js";

describe("protocol type guards", () => {
  it("isAgentMessage accepts valid heartbeat", () => {
    expect(
      isAgentMessage({
        type: "heartbeat",
        hostname: "worker-1",
        uptime: 100,
        version: "0.1.0",
        memUsage: 1024,
        cpuUsage: 0.5,
      }),
    ).toBe(true);
  });

  it("isAgentMessage accepts valid exec-exit", () => {
    expect(
      isAgentMessage({ type: "exec-exit", requestId: "abc", exitCode: 0 }),
    ).toBe(true);
  });

  it("isAgentMessage rejects unknown type", () => {
    expect(isAgentMessage({ type: "unknown-type" })).toBe(false);
  });

  it("isAgentMessage rejects non-object", () => {
    expect(isAgentMessage("heartbeat")).toBe(false);
    expect(isAgentMessage(null)).toBe(false);
    expect(isAgentMessage(42)).toBe(false);
  });

  it("isServerMessage accepts valid exec", () => {
    expect(
      isServerMessage({
        type: "exec",
        requestId: "abc",
        command: "ls",
        args: ["-la"],
        timeout: 30000,
        tty: false,
      }),
    ).toBe(true);
  });

  it("isServerMessage accepts server-shutdown", () => {
    expect(
      isServerMessage({ type: "server-shutdown", reconnectAfter: 5000 }),
    ).toBe(true);
  });

  it("isServerMessage rejects agent message types", () => {
    expect(isServerMessage({ type: "heartbeat" })).toBe(false);
  });
});

describe("parseAgentMessage", () => {
  it("parses valid JSON", () => {
    const msg = parseAgentMessage(
      JSON.stringify({ type: "heartbeat", hostname: "w1", uptime: 1, version: "0.1.0", memUsage: 0, cpuUsage: 0 }),
    );
    expect(msg.type).toBe("heartbeat");
  });

  it("throws on invalid JSON", () => {
    expect(() => parseAgentMessage("not json")).toThrow();
  });

  it("throws on invalid message type", () => {
    expect(() => parseAgentMessage(JSON.stringify({ type: "bogus" }))).toThrow(
      "Invalid agent message",
    );
  });
});

describe("parseServerMessage", () => {
  it("parses valid JSON", () => {
    const msg = parseServerMessage(
      JSON.stringify({ type: "heartbeat-ack", serverTime: "2026-01-01T00:00:00Z" }),
    );
    expect(msg.type).toBe("heartbeat-ack");
  });

  it("throws on agent message type", () => {
    expect(() =>
      parseServerMessage(JSON.stringify({ type: "heartbeat" })),
    ).toThrow("Invalid server message");
  });
});

describe("generateRequestId", () => {
  it("returns a string", () => {
    expect(typeof generateRequestId()).toBe("string");
  });

  it("returns unique IDs", () => {
    const ids = new Set(Array.from({ length: 100 }, () => generateRequestId()));
    expect(ids.size).toBe(100);
  });
});