MCMapper-Backend/api/src/ws-gateway.test.ts

import { describe, test, expect, beforeAll, afterAll } from "bun:test";
import { and, eq } from "drizzle-orm";
import { db } from "./db/client";
import { chatMessages, chunkColumns, chunkSections } from "./db/schema";
import { redis, DIRTY_CHUNK_STREAM } from "./redis";
import { wsGateway, chunkOf, getModSocket } from "./ws-gateway";
import { createTestServer, deleteTestServer, FakeSocket } from "./test-helpers";
import { redeemLinkCode } from "./link";
import { chatChannel } from "./chat";

describe("chunkOf", () => {
  test("floors toward negative infinity, matching Math.floor(coord/16)", () => {
    expect(chunkOf(0)).toBe(0);
    expect(chunkOf(15)).toBe(0);
    expect(chunkOf(16)).toBe(1);
    expect(chunkOf(-1)).toBe(-1);
    expect(chunkOf(-16)).toBe(-1);
    expect(chunkOf(-17)).toBe(-2);
  });
});

describe("wsGateway.message", () => {
  let server: { id: string; token: string };

  beforeAll(async () => {
    server = await createTestServer("ws-gateway");
  });

  afterAll(async () => {
    await deleteTestServer(server.id);
  });

  test("hello with a valid token acknowledges with the server id", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));
    expect(socket.closed).toBe(false);
    expect(socket.lastMessage()).toEqual({ type: "hello_ack", ok: true, serverId: server.id });
  });

  test("hello with an invalid token is rejected and the connection closed", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: "not-a-real-token" }));
    expect(socket.closed).toBe(true);
    expect(socket.lastMessage()).toEqual({ type: "hello_ack", ok: false, error: "invalid_token" });
  });

  test("a message before hello is rejected as not_authenticated", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "columns", dimension: 0, columns: [] }));
    expect(socket.closed).toBe(true);
    expect(socket.lastMessage()).toEqual({ type: "error", error: "not_authenticated" });
  });

  test("invalid JSON gets an error reply, not a thrown exception", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, "{not json");
    expect(socket.lastMessage()).toEqual({ type: "error", error: "invalid_json" });
  });

  test("columns upserts chunk_columns and marks exactly the touched chunks dirty", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));

    const streamLenBefore = await redis.xlen(DIRTY_CHUNK_STREAM);

    // Two columns in chunk (0,0), one in chunk (1,0) — should dedupe to exactly 2 dirty events.
    await wsGateway.message(
      socket,
      JSON.stringify({
        type: "columns",
        dimension: 0,
        columns: [
          { x: 1, z: 1, height: 64, blockId: 2, blockMeta: 0 },
          { x: 2, z: 2, height: 65, blockId: 3, blockMeta: 0 },
          { x: 16, z: 1, height: 70, blockId: 1, blockMeta: 0 },
        ],
      }),
    );

    const rows = await db
      .select()
      .from(chunkColumns)
      .where(and(eq(chunkColumns.serverId, server.id), eq(chunkColumns.dimension, 0)));
    expect(rows).toHaveLength(3);
    const byX = Object.fromEntries(rows.map((r) => [r.x, r]));
    expect(byX[1]!.blockId).toBe(2);
    expect(byX[1]!.height).toBe(64);
    expect(byX[16]!.blockId).toBe(1);

    const streamLenAfter = await redis.xlen(DIRTY_CHUNK_STREAM);
    expect(streamLenAfter - streamLenBefore).toBe(2);
  });

  test("a repeated column overwrites rather than duplicating", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));

    const send = (x: number, blockId: number) =>
      wsGateway.message(
        socket,
        JSON.stringify({
          type: "columns",
          dimension: 0,
          columns: [{ x, z: 100, height: 64, blockId, blockMeta: 0 }],
        }),
      );

    await send(50, 2);
    await send(50, 3);

    const rows = await db
      .select()
      .from(chunkColumns)
      .where(and(eq(chunkColumns.serverId, server.id), eq(chunkColumns.x, 50), eq(chunkColumns.z, 100)));
    expect(rows).toHaveLength(1);
    expect(rows[0]!.blockId).toBe(3);
  });

  test("sections upserts chunk_sections and marks the chunk dirty", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));

    const blocks = Buffer.alloc(8192).toString("base64"); // all-air section is fine for this test
    const streamLenBefore = await redis.xlen(DIRTY_CHUNK_STREAM);

    await wsGateway.message(
      socket,
      JSON.stringify({
        type: "sections",
        dimension: 0,
        chunkX: 9,
        chunkZ: -3,
        sections: [{ sectionY: 4, blocks }],
      }),
    );

    const rows = await db
      .select()
      .from(chunkSections)
      .where(
        and(
          eq(chunkSections.serverId, server.id),
          eq(chunkSections.x, 9),
          eq(chunkSections.z, -3),
          eq(chunkSections.sectionY, 4),
        ),
      );
    expect(rows).toHaveLength(1);
    expect(rows[0]!.blocks).toBe(blocks);

    const streamLenAfter = await redis.xlen(DIRTY_CHUNK_STREAM);
    expect(streamLenAfter - streamLenBefore).toBe(1);
  });

  test("close() forgets the connection's authentication state", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));
    wsGateway.close(socket);

    const socket2 = new FakeSocket();
    socket2.id = socket.id; // simulate the same connection id being reused post-close
    await wsGateway.message(socket2, JSON.stringify({ type: "columns", dimension: 0, columns: [] }));
    expect(socket2.lastMessage()).toEqual({ type: "error", error: "not_authenticated" });
  });

  test("link_request stores a redeemable code carrying the mod-supplied identity", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));

    const code = crypto.randomUUID().slice(0, 6);
    await wsGateway.message(
      socket,
      JSON.stringify({
        type: "link_request",
        code,
        uuid: "mod-supplied-uuid",
        username: "ModPlayer",
        authMode: "online",
      }),
    );

    const redeemed = await redeemLinkCode(code);
    expect(redeemed.ok).toBe(true);
    if (!redeemed.ok) return;
    expect(redeemed.account.username).toBe("ModPlayer");
  });

  test("getModSocket returns the authenticated connection for a server, and forgets it on close", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));
    expect(getModSocket(server.id)).toBe(socket);

    wsGateway.close(socket);
    expect(getModSocket(server.id)).toBeUndefined();
  });

  test("chat from the mod persists it and publishes it for web viewers", async () => {
    const socket = new FakeSocket();
    await wsGateway.message(socket, JSON.stringify({ type: "hello", token: server.token }));

    const sub = redis.duplicate();
    await sub.subscribe(chatChannel(server.id));
    const received = new Promise<string>((resolve) => {
      sub.once("message", (_channel, message) => resolve(message));
    });

    await wsGateway.message(
      socket,
      JSON.stringify({ type: "chat", uuid: "chatter-uuid", username: "InGamePlayer", message: "gg" }),
    );

    expect(JSON.parse(await received)).toEqual({ username: "InGamePlayer", message: "gg", source: "game" });
    sub.disconnect();

    const rows = await db
      .select()
      .from(chatMessages)
      .where(and(eq(chatMessages.serverId, server.id), eq(chatMessages.message, "gg")));
    expect(rows).toHaveLength(1);
    expect(rows[0]!.source).toBe("game");
  });
});