matrix-backend/tests/websocket.test.ts

import { describe, it, expect, vi, beforeEach, type Mocked } from "vitest";
import { Server } from "http";
import { WebSocket, Server as WebSocketServer } from "ws";
import { ExtendedWebSocketServer } from "../src/websocket";
import { WebsocketServerEventHandler } from "../src/utils/websocket/websocketServerEventHandler";
import { WebsocketEventHandler } from "../src/utils/websocket/websocketEventHandler";
import { getEventListeners } from "../src/utils/websocket/websocketCustomEvents/websocketEventUtils";
import { createMockUserService} from "./helpers/testSetup";
import {UserService} from "../src/services/db/UserService";
import {SpotifyPollingService} from "../src/services/spotifyPollingService";
import { USER_UPDATED_EVENT, SPOTIFY_STATE_UPDATED_EVENT } from "../src/utils/eventBus";
import { WebsocketEventType } from "../src/utils/websocket/websocketCustomEvents/websocketEventType";
import {WeatherPollingService} from "../src/services/weatherPollingService";

let mockWssInstance: Mocked<WebSocketServer>;
let mockServerEventHandler: Mocked<WebsocketServerEventHandler>;

const eventBusListeners = new Map<string, (...args: any[]) => void>();

vi.mock("../src/utils/eventBus", () => ({
    appEventBus: {
        on: vi.fn((event, listener) => {
            eventBusListeners.set(event, listener);
        }),
        emit: vi.fn(),
    },
    USER_UPDATED_EVENT: 'user:updated',
    SPOTIFY_STATE_UPDATED_EVENT: 'spotify:state-updated',
    WEATHER_STATE_UPDATED_EVENT: 'weather-state:updated',
}));

vi.mock("ws", () => ({
    Server: vi.fn().mockImplementation(() => mockWssInstance),
    WebSocket: { OPEN: 1, CLOSED: 3 },
}));

vi.mock("../src/utils/verifyClient");
vi.mock("../src/utils/websocket/websocketServerEventHandler", () => ({
    WebsocketServerEventHandler: vi.fn().mockImplementation(() => mockServerEventHandler),
}));
vi.mock("../src/utils/websocket/websocketEventHandler");
vi.mock("../src/utils/websocket/websocketCustomEvents/websocketEventUtils");

describe("ExtendedWebSocketServer", () => {
    let mockHttpServer: Mocked<Server>;
    let extendedWss: ExtendedWebSocketServer;
    let mockUserService: Mocked<UserService>;
    let mockSpotifyPollingService: Mocked<SpotifyPollingService>
    let mockWeatherPollingService: Mocked<WeatherPollingService>;

    beforeEach(() => {
        vi.clearAllMocks();

        mockHttpServer = {} as Mocked<Server>;

        mockServerEventHandler = {
            enableConnectionEvent: vi.fn(),
            enableHeartbeat: vi.fn(),
            enableCloseEvent: vi.fn(),
        } as unknown as Mocked<WebsocketServerEventHandler>;

        mockWssInstance = {
            clients: new Set(),
            on: vi.fn(),
            close: vi.fn(),
        } as unknown as Mocked<WebSocketServer>;

        mockSpotifyPollingService = {} as any;
        mockWeatherPollingService = {} as any;

        mockUserService = createMockUserService();

        extendedWss = new ExtendedWebSocketServer(mockHttpServer, mockUserService, mockSpotifyPollingService, mockWeatherPollingService);
    });

    describe("Constructor and Setup", () => {
        it("should create a new WebSocket.Server", () => {
            expect(WebSocketServer).toHaveBeenCalledWith({
                server: mockHttpServer,
                verifyClient: expect.any(Function),
            });
        });

        it("should create and use a WebsocketServerEventHandler with the correct service", () => {
            expect(WebsocketServerEventHandler).toHaveBeenCalledWith(mockWssInstance, mockUserService);
        });

        it("should enable the heartbeat", () => {
            expect(mockServerEventHandler.enableHeartbeat).toHaveBeenCalledWith(30000);
        });

        it("should register a connection handler", () => {
            expect(mockServerEventHandler.enableConnectionEvent).toHaveBeenCalledWith(expect.any(Function));
        });
    });

    describe("broadcast", () => {
        it("should send a message to all connected clients that are OPEN", () => {
            const client1 = { readyState: WebSocket.OPEN, send: vi.fn() };
            const client2 = { readyState: WebSocket.CLOSED, send: vi.fn() };
            mockWssInstance.clients.add(client1 as any).add(client2 as any);
            extendedWss.broadcast("hello");
            expect(client1.send).toHaveBeenCalledWith("hello", { binary: false });
            expect(client2.send).not.toHaveBeenCalled();
        });
    });

    describe("sendMessageToUser", () => {
        it("should send a message to a specific user by their UUID", () => {
            const client1 = { readyState: WebSocket.OPEN, payload: { uuid: "uuid-1" }, send: vi.fn() };
            const client2 = { readyState: WebSocket.OPEN, payload: { uuid: "uuid-2" }, send: vi.fn() };
            mockWssInstance.clients.add(client1 as any).add(client2 as any);
            extendedWss.sendMessageToUser("uuid-1", "private");
            expect(client1.send).toHaveBeenCalledWith("private", { binary: false });
            expect(client2.send).not.toHaveBeenCalled();
        });
    });

    describe("Connection Handler Logic", () => {
        let connectionHandler: (ws: any, req: any) => void;
        let mockWsClient: any;
        let mockClientEventHandler: Mocked<WebsocketEventHandler>;

        beforeEach(() => {
            connectionHandler = vi.mocked(mockServerEventHandler.enableConnectionEvent).mock.calls[0][0];
            mockWsClient = {
                emit: vi.fn(), on: vi.fn(), user: { lastState: { global: { mode: "idle" } } },
            };
            mockClientEventHandler = {
                enableErrorEvent: vi.fn(),
                enablePongEvent: vi.fn(),
                enableMessageEvent: vi.fn(),
                enableDisconnectEvent: vi.fn(),
                registerCustomEvents: vi.fn(),
            } as unknown as Mocked<WebsocketEventHandler>;

            vi.mocked(WebsocketEventHandler).mockImplementation(() => mockClientEventHandler);
            vi.mocked(getEventListeners).mockReturnValue([{ event: "custom", handler: vi.fn() } as any]);
        });

        it("should create and configure a WebsocketEventHandler for new clients", () => {
            connectionHandler(mockWsClient, {});
            expect(vi.mocked(WebsocketEventHandler)).toHaveBeenCalledWith(mockWsClient, mockSpotifyPollingService, mockWeatherPollingService);
            expect(mockClientEventHandler.enableErrorEvent).toHaveBeenCalled();
            expect(mockClientEventHandler.enablePongEvent).toHaveBeenCalled();
            expect(mockClientEventHandler.enableMessageEvent).toHaveBeenCalled();
            expect(mockClientEventHandler.enableDisconnectEvent).toHaveBeenCalled();
            expect(mockClientEventHandler.registerCustomEvents).toHaveBeenCalled();
        });

        it("should emit initial events to the new client", () => {
            connectionHandler(mockWsClient, {});
            expect(mockWsClient.emit).toHaveBeenCalledWith("GET_STATE", {});
            expect(mockWsClient.emit).toHaveBeenCalledWith("GET_SETTINGS", {});
        });
    });

    describe("_listenForAppEvents", () => {
        let mockClient: any;

        beforeEach(() => {
            mockClient = {
                readyState: WebSocket.OPEN,
                payload: { uuid: "user-123" },
                send: vi.fn(),
                emit: vi.fn(),
            };
            mockWssInstance.clients.add(mockClient);
        });

        it("should listen for USER_UPDATED_EVENT and emit to the correct client", () => {
            const userUpdateListener = eventBusListeners.get(USER_UPDATED_EVENT);
            expect(userUpdateListener).toBeDefined();

            const updatedUserPayload = { uuid: "user-123", name: "Neuer Name" };

            userUpdateListener!(updatedUserPayload);

            expect(mockClient.emit).toHaveBeenCalledOnce();
            expect(mockClient.emit).toHaveBeenCalledWith(
                WebsocketEventType.UPDATE_USER_SINGLE,
                updatedUserPayload
            );
        });

        it("should listen for SPOTIFY_STATE_UPDATED_EVENT and send to the correct client", () => {
            const spotifyStateListener = eventBusListeners.get(SPOTIFY_STATE_UPDATED_EVENT);
            expect(spotifyStateListener).toBeDefined();

            const spotifyUpdatePayload = { state: { item: { name: "Neuer Song" } } };
            const eventPayload = { uuid: "user-123", ...spotifyUpdatePayload };

            spotifyStateListener!(eventPayload);

            expect(mockClient.send).toHaveBeenCalledOnce();
            const expectedMessage = JSON.stringify({
                type: "SPOTIFY_UPDATE",
                payload: spotifyUpdatePayload.state,
            });
            expect(mockClient.send).toHaveBeenCalledWith(expectedMessage, { binary: false });
        });

        it("should not send a message if the target client is not connected", () => {
            const userUpdateListener = eventBusListeners.get(USER_UPDATED_EVENT);
            const spotifyStateListener = eventBusListeners.get(SPOTIFY_STATE_UPDATED_EVENT);

            const eventPayload = { uuid: "user-unknown", name: "some data" };

            userUpdateListener!(eventPayload);

            spotifyStateListener!(eventPayload);

            expect(mockClient.send).not.toHaveBeenCalled();
            expect(mockClient.emit).not.toHaveBeenCalled();
        });
    });
});