proto-socket/typescript/test/ws.test.ts

import * as fs from "node:fs";
import * as path from "node:path";
import { fileURLToPath } from "node:url";

import { create } from "@bufbuild/protobuf";
import { afterEach, describe, expect, test } from "vitest";

import {
  addListenerTyped,
  addRequestListenerTyped,
  parserFromSchema,
  sendRequestTyped,
} from "../src/communicator.js";
import { TestDataSchema, type TestData } from "../src/packets/message_common_pb.js";
import { connectWs, connectWss, WsClient } from "../src/ws_client.js";
import { WsServer } from "../src/ws_server.js";

const certsDir = path.join(path.dirname(fileURLToPath(import.meta.url)), "certs");
const cert = fs.readFileSync(path.join(certsDir, "server.crt"));
const key = fs.readFileSync(path.join(certsDir, "server.key"));

function parserMap() {
  return new Map([[TestDataSchema.typeName, parserFromSchema(TestDataSchema)]]);
}

async function waitForClientMessage(client: WsClient): Promise<TestData> {
  return new Promise<TestData>((resolve) => {
    addListenerTyped(client.communicator, TestDataSchema, (msg) => resolve(msg));
  });
}

describe("WS", () => {
  const cleanup: Array<() => Promise<void>> = [];

  afterEach(async () => {
    while (cleanup.length > 0) {
      const fn = cleanup.pop();
      if (fn) {
        await fn();
      }
    }
  });

  test("connectWs sends and receives TestData", async () => {
    const server = new WsServer("127.0.0.1", 0, "/", (ws) => new WsClient(ws, 0, 0, parserMap()));
    cleanup.push(async () => server.stop());
    await server.start();

    server.onClientConnected = (client) => {
      addListenerTyped(client.communicator, TestDataSchema, async (msg) => {
        if (msg.index === 11 && msg.message === "hello over ws") {
          await client.communicator.send(
            create(TestDataSchema, { index: 200, message: "push from ws server" }),
          );
        }
      });
    };

    const client = await connectWs("127.0.0.1", server.port, "/", 0, 0, parserMap());
    cleanup.push(async () => client.close());

    const pushed = waitForClientMessage(client);
    await client.communicator.send(create(TestDataSchema, { index: 11, message: "hello over ws" }));

    await expect(pushed).resolves.toMatchObject({
      index: 200,
      message: "push from ws server",
    });
  });

  test("sendRequest/response roundtrip over WS", async () => {
    const server = new WsServer("127.0.0.1", 0, "/", (ws) => new WsClient(ws, 0, 0, parserMap()));
    cleanup.push(async () => server.stop());
    await server.start();

    server.onClientConnected = (client) => {
      addRequestListenerTyped(client.communicator, TestDataSchema, (req) =>
        create(TestDataSchema, {
          index: req.index * 2,
          message: `echo: ${req.message}`,
        }),
      );
    };

    const client = await connectWs("127.0.0.1", server.port, "/", 0, 0, parserMap());
    cleanup.push(async () => client.close());

    await expect(
      sendRequestTyped(
        client.communicator,
        create(TestDataSchema, { index: 21, message: "request over ws" }),
        TestDataSchema,
        500,
      ),
    ).resolves.toMatchObject({
      index: 42,
      message: "echo: request over ws",
    });
  });

  test("connectWss sends and receives TestData", async () => {
    const server = new WsServer("127.0.0.1", 0, "/", (ws) => new WsClient(ws, 0, 0, parserMap()), {
      cert,
      key,
    });
    cleanup.push(async () => server.stop());
    await server.start();

    server.onClientConnected = (client) => {
      addListenerTyped(client.communicator, TestDataSchema, async (msg) => {
        if (msg.index === 12 && msg.message === "hello over wss") {
          await client.communicator.send(
            create(TestDataSchema, { index: 201, message: "push from wss server" }),
          );
        }
      });
    };

    const client = await connectWss(
      "127.0.0.1",
      server.port,
      "/",
      { ca: cert },
      0,
      0,
      parserMap(),
    );
    cleanup.push(async () => client.close());

    const pushed = waitForClientMessage(client);
    await client.communicator.send(create(TestDataSchema, { index: 12, message: "hello over wss" }));

    await expect(pushed).resolves.toMatchObject({
      index: 201,
      message: "push from wss server",
    });
  });

  test("WSS sendRequest/response roundtrip", async () => {
    const server = new WsServer("127.0.0.1", 0, "/", (ws) => new WsClient(ws, 0, 0, parserMap()), {
      cert,
      key,
    });
    cleanup.push(async () => server.stop());
    await server.start();

    server.onClientConnected = (client) => {
      addRequestListenerTyped(client.communicator, TestDataSchema, (req) =>
        create(TestDataSchema, {
          index: req.index * 2,
          message: `wss echo: ${req.message}`,
        }),
      );
    };

    const client = await connectWss(
      "127.0.0.1",
      server.port,
      "/",
      { ca: cert },
      0,
      0,
      parserMap(),
    );
    cleanup.push(async () => client.close());

    await expect(
      sendRequestTyped(
        client.communicator,
        create(TestDataSchema, { index: 22, message: "request over wss" }),
        TestDataSchema,
        500,
      ),
    ).resolves.toMatchObject({
      index: 44,
      message: "wss echo: request over wss",
    });
  });

  test("concurrent sendRequest/response roundtrip over WS", async () => {
    const server = new WsServer("127.0.0.1", 0, "/", (ws) => new WsClient(ws, 0, 0, parserMap()));
    cleanup.push(async () => server.stop());
    await server.start();

    server.onClientConnected = (client) => {
      addRequestListenerTyped(client.communicator, TestDataSchema, (req) =>
        create(TestDataSchema, {
          index: req.index * 2,
          message: `echo: ${req.message}`,
        }),
      );
    };

    const client = await connectWs("127.0.0.1", server.port, "/", 0, 0, parserMap());
    cleanup.push(async () => client.close());

    const results = await Promise.all(
      Array.from({ length: 5 }, (_, i) =>
        sendRequestTyped(
          client.communicator,
          create(TestDataSchema, { index: 30 + i, message: `request ${i}` }),
          TestDataSchema,
          2000,
        ),
      ),
    );

    results.forEach((res, i) => {
      expect(res.index).toBe((30 + i) * 2);
      expect(res.message).toBe(`echo: request ${i}`);
    });
  });
});