clientmanager: add some unit tests

server/balancer.ts

@@ -35,7 +35,7 @@ export function initBalancerConnections() {
 	});
 	wss.on("connection", ws => {
 		log.debug("New balancer connection");
-		const conn = new BalancerConnection(ws);
+		const conn = new BalancerConnectionReal(ws);
 		balancerManager.addBalancerConnection(conn);
 	});
 	wss.on("error", error => {
@@ -115,8 +115,8 @@ class BalancerManager {
 
 export const balancerManager = new BalancerManager();
 
-type BalancerManagerEvemts = BalancerConnectionEvents;
-type BalancerManagerEventHandlers<E> = E extends "connect"
+export type BalancerManagerEvemts = BalancerConnectionEvents;
+export type BalancerManagerEventHandlers<E> = E extends "connect"
 	? (conn: BalancerConnection) => void
 	: E extends "disconnect"
 	? (conn: BalancerConnection) => void
@@ -126,8 +126,8 @@ type BalancerManagerEventHandlers<E> = E extends "connect"
 	? (conn: BalancerConnection, error: WebSocket.ErrorEvent) => void
 	: never;
 
-type BalancerConnectionEvents = "connect" | "disconnect" | "message" | "error";
-type BalancerConnectionEventHandlers<E> = E extends "connect"
+export type BalancerConnectionEvents = "connect" | "disconnect" | "message" | "error";
+export type BalancerConnectionEventHandlers<E> = E extends "connect"
 	? () => void
 	: E extends "disconnect"
 	? (code: number, reason: string) => void
@@ -137,15 +137,35 @@ type BalancerConnectionEventHandlers<E> = E extends "connect"
 	? (error: WebSocket.ErrorEvent) => void
 	: never;
 
-/** Manages the websocket connection to a Balancer. */
-export class BalancerConnection {
+export abstract class BalancerConnection {
 	/** A local identifier for the balancer. Other monoliths will have different IDs for the same balancer. */
 	id: string;
+	protected bus: EventEmitter = new EventEmitter();
+
+	constructor() {
+		this.id = uuidv4();
+	}
+
+	protected emit<E extends BalancerConnectionEvents>(
+		event: E,
+		...args: Parameters<BalancerConnectionEventHandlers<E>>
+	) {
+		this.bus.emit(event, ...args);
+	}
+
+	on<E extends BalancerConnectionEvents>(event: E, handler: BalancerConnectionEventHandlers<E>) {
+		this.bus.on(event, handler);
+	}
+
+	abstract send(message: MsgM2B): Result<void, Error>;
+}
+
+/** Manages the websocket connection to a Balancer. */
+export class BalancerConnectionReal extends BalancerConnection {
 	private socket: WebSocket;
-	private bus: EventEmitter = new EventEmitter();
 
 	constructor(socket: WebSocket) {
-		this.id = uuidv4();
+		super();
 		this.socket = socket;
 
 		this.socket.on("open", this.onSocketConnect.bind(this));
@@ -201,17 +221,6 @@ export class BalancerConnection {
 		this.emit("error", event);
 	}
 
-	private emit<E extends BalancerConnectionEvents>(
-		event: E,
-		...args: Parameters<BalancerConnectionEventHandlers<E>>
-	) {
-		this.bus.emit(event, ...args);
-	}
-
-	on<E extends BalancerConnectionEvents>(event: E, handler: BalancerConnectionEventHandlers<E>) {
-		this.bus.on(event, handler);
-	}
-
 	send(message: MsgM2B): Result<void, Error> {
 		if (this.socket === null) {
 			return err(new Error("Not connected"));

server/client.ts

@@ -6,7 +6,7 @@ import { v4 as uuidv4 } from "uuid";
 import EventEmitter from "events";
 import { getLogger } from "./logger";
 import { getSessionInfo } from "./auth/tokens";
-import { BalancerConnection } from "./balancer";
+import { BalancerConnection, BalancerConnectionReal } from "./balancer";
 import { replacer } from "../common/serialize";
 
 const log = getLogger("client");

server/clientmanager.ts

@@ -55,10 +55,12 @@ export async function setup(): Promise<void> {
 	balancerManager.on("error", onBalancerError);
 	initBalancerConnections();
 
-	log.silly("creating redis subscriber");
-	const redisSubscriber = await createSubscriber();
-	log.silly("subscribing to announcement channel");
-	await redisSubscriber.subscribe(ANNOUNCEMENT_CHANNEL, onAnnouncement);
+	if (conf.get("env") !== "test") {
+		log.silly("creating redis subscriber");
+		const redisSubscriber = await createSubscriber();
+		log.silly("subscribing to announcement channel");
+		await redisSubscriber.subscribe(ANNOUNCEMENT_CHANNEL, onAnnouncement);
+	}
 }
 
 /**
@@ -69,6 +71,10 @@ async function onDirectConnect(socket: WebSocket, req: express.Request) {
 	const roomName = req.url.split("/").slice(-1)[0];
 	log.debug(`connection received: ${roomName}, waiting for auth token...`);
 	const client = new DirectClient(roomName, socket);
+	addClient(client);
+}
+
+export function addClient(client: Client) {
 	connections.push(client);
 	client.on("auth", onClientAuth);
 	client.on("message", onClientMessage);
@@ -427,6 +433,10 @@ function getClient(id: ClientId): Client | undefined {
 	return undefined;
 }
 
+function getClientsInRoom(roomName: string): Client[] {
+	return roomJoins.get(roomName) ?? [];
+}
+
 setInterval(() => {
 	for (const client of connections) {
 		if (client instanceof DirectClient) {
@@ -472,4 +482,7 @@ export default {
 	onUserModified,
 	getClientByToken,
 	makeRoomRequest,
+	addClient,
+	getClient,
+	getClientsInRoom,
 };

server/tests/unit/clientmanager.spec.ts

@@ -0,0 +1,120 @@
+import clientmanager from "../../clientmanager";
+import {
+	BalancerConnection,
+	BalancerConnectionEventHandlers,
+	BalancerConnectionEvents,
+	BalancerConnectionReal,
+	MsgM2B,
+	balancerManager,
+} from "../../balancer";
+import { BalancerClient, Client } from "../../client";
+import { OttWebsocketError } from "common/models/types";
+import { buildClients } from "../../redisclient";
+import { Result, ok } from "../../../common/result";
+import roommanager from "../../roommanager";
+import { loadModels } from "../../models";
+
+class TestClient extends Client {
+	sendRawMock = jest.fn();
+	kickMock = jest.fn();
+
+	sendRaw(msg: string): void {
+		this.sendRawMock(msg);
+	}
+
+	kick(code: OttWebsocketError): void {
+		this.kickMock(code);
+		this.emit("disconnect", this);
+	}
+}
+
+class BalancerConnectionMock extends BalancerConnection {
+	sendMock = jest.fn();
+
+	constructor() {
+		super();
+	}
+
+	send(msg: MsgM2B): Result<void, Error> {
+		this.sendMock(msg);
+		return ok(undefined);
+	}
+
+	public emit<E extends BalancerConnectionEvents>(
+		event: E,
+		...args: Parameters<BalancerConnectionEventHandlers<E>>
+	) {
+		super.emit(event, ...args);
+	}
+}
+
+describe("ClientManager", () => {
+	beforeAll(async () => {
+		loadModels();
+		await buildClients();
+		await clientmanager.setup();
+	});
+
+	beforeEach(async () => {
+		await roommanager.createRoom({
+			name: "foo",
+			isTemporary: true,
+		});
+	});
+
+	afterEach(async () => {
+		await roommanager.unloadRoom("foo");
+	});
+
+	it("should add clients", () => {
+		const client = new TestClient("foo");
+		clientmanager.addClient(client);
+		expect(clientmanager.getClient(client.id)).toBe(client);
+	});
+
+	it("should add clients to roomJoins when they auth", async () => {
+		const client = new TestClient("foo");
+		clientmanager.addClient(client);
+		client.emit("auth", client, "token", { isLoggedIn: false, username: "foo" });
+		await new Promise(resolve => setTimeout(resolve, 100));
+		const joins = clientmanager.getClientsInRoom("foo");
+		expect(joins).toHaveLength(1);
+	});
+
+	it("should remove clients when they disconnect", () => {
+		const client = new TestClient("foo");
+		clientmanager.addClient(client);
+		client.emit("disconnect", client);
+		expect(clientmanager.getClient(client.id)).toBeUndefined();
+	});
+
+	it("should remove clients from roomJoins when they disconnect", async () => {
+		const client = new TestClient("foo");
+		clientmanager.addClient(client);
+		client.emit("auth", client, "token", { isLoggedIn: false, username: "foo" });
+		await new Promise(resolve => setTimeout(resolve, 100));
+		const joins = clientmanager.getClientsInRoom("foo");
+		expect(joins).toHaveLength(1);
+
+		client.emit("disconnect", client);
+		const joins2 = clientmanager.getClientsInRoom("foo");
+		expect(joins2).toHaveLength(0);
+	});
+
+	it("should disconnect all clients when a balancer disconnects", async () => {
+		const mockBalancerCon = new BalancerConnectionMock();
+		balancerManager.addBalancerConnection(mockBalancerCon);
+		const client = new BalancerClient("foo", "foo", mockBalancerCon);
+		clientmanager.addClient(client);
+		client.emit("auth", client, "token", { isLoggedIn: false, username: "foo" });
+		await new Promise(resolve => setTimeout(resolve, 100));
+		const joins = clientmanager.getClientsInRoom("foo");
+		expect(joins).toHaveLength(1);
+
+		mockBalancerCon.emit("disconnect", 1000, "reason");
+
+		expect(clientmanager.getClient(client.id)).toBeUndefined();
+		const joins2 = clientmanager.getClientsInRoom("foo");
+		expect(joins2).toHaveLength(0);
+	});
+});