diff --git a/packages/mesh/src/server/index.ts b/packages/mesh/src/server/index.ts index 2aa977b..8c820ef 100644 --- a/packages/mesh/src/server/index.ts +++ b/packages/mesh/src/server/index.ts @@ -1089,4 +1089,30 @@ export class MeshServer extends WebSocketServer { callback(); } } + + /** + * Registers a callback function to be executed when a new connection is established. + * + * @param {(connection: Connection) => Promise | void} callback - The function to execute when a new connection is established. + * @returns {MeshServer} The server instance for method chaining. + */ + onConnection( + callback: (connection: Connection) => Promise | void + ): MeshServer { + this.on("connected", callback); + return this; + } + + /** + * Registers a callback function to be executed when a connection is closed. + * + * @param {(connection: Connection) => Promise | void} callback - The function to execute when a connection is closed. + * @returns {MeshServer} The server instance for method chaining. + */ + onDisconnection( + callback: (connection: Connection) => Promise | void + ): MeshServer { + this.on("disconnected", callback); + return this; + } } diff --git a/packages/mesh/src/tests/basic.test.ts b/packages/mesh/src/tests/basic.test.ts index 0e200b8..29e5ee0 100644 --- a/packages/mesh/src/tests/basic.test.ts +++ b/packages/mesh/src/tests/basic.test.ts @@ -23,7 +23,9 @@ const flushRedis = async () => { await redis.quit(); }; -describe("KeepAliveServer", () => { +const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +describe("MeshServer", () => { const port = 8126; let server: MeshServer; let clientA: MeshClient; @@ -129,5 +131,58 @@ describe("KeepAliveServer", () => { await server.connectionManager.getAllMetadataForRoom("room-b"); expect(roomBMetadata).toEqual([{ [connectionB.id]: metadataB }]); }); + + test("onConnection callback is executed when a client connects", async () => { + let connectionReceived: any = null; + const connectionPromise = new Promise((resolve) => { + server.onConnection((connection) => { + connectionReceived = connection; + resolve(); + }); + }); + + await clientA.connect(); + await connectionPromise; + + expect(connectionReceived).not.toBeNull(); + + if (!connectionReceived) { + return; + } + + expect(connectionReceived.id).toBeDefined(); + expect(connectionReceived.isDead).toBe(false); + + const connections = server.connectionManager.getLocalConnections(); + expect(connections).toContain(connectionReceived); + }); + + test("onDisconnection callback is executed when a client disconnects", async () => { + let disconnectedConnection: any = null; + const disconnectionPromise = new Promise((resolve) => { + server.onDisconnection((connection) => { + disconnectedConnection = connection; + resolve(); + }); + }); + + await clientA.connect(); + await wait(100); + const connections = server.connectionManager.getLocalConnections(); + const connectionBeforeDisconnect = connections[0]; + + expect(connectionBeforeDisconnect).toBeDefined(); + const connectionId = connectionBeforeDisconnect?.id; + + await clientA.close(); + await disconnectionPromise; + + expect(disconnectedConnection).not.toBeNull(); + + if (disconnectedConnection && connectionId) { + expect(disconnectedConnection.id).toBe(connectionId); + expect(disconnectedConnection.isDead).toBe(true); + } + }); }); }); diff --git a/packages/mesh/src/tests/client.test.ts b/packages/mesh/src/tests/client.test.ts index c27ace3..6af336b 100644 --- a/packages/mesh/src/tests/client.test.ts +++ b/packages/mesh/src/tests/client.test.ts @@ -23,7 +23,7 @@ const flushRedis = async () => { await redis.quit(); }; -describe("KeepAliveClient", () => { +describe("MeshClient", () => { const port = 8127; let server: MeshServer; let client: MeshClient; diff --git a/packages/mesh/src/tests/rooms.test.ts b/packages/mesh/src/tests/rooms.test.ts index 028f829..8f190f3 100644 --- a/packages/mesh/src/tests/rooms.test.ts +++ b/packages/mesh/src/tests/rooms.test.ts @@ -23,7 +23,7 @@ const flushRedis = async () => { await redis.quit(); }; -describe("KeepAliveServer", () => { +describe("MeshServer", () => { const port = 8128; let server: MeshServer; let clientA: MeshClient;