From 94a725a717e8b4ba9dd4d2a4dcd87f4192a0828e Mon Sep 17 00:00:00 2001
From: Gus Narea <gnarea@users.noreply.github.com>
Date: Wed, 28 Jul 2021 19:06:00 +0100
Subject: [PATCH] feat(MockClient): Implement method to run callback whilst
 connected to server

---
 src/lib/MockClient.ts | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/lib/MockClient.ts b/src/lib/MockClient.ts
index 4e7e860..12bb7dd 100644
--- a/src/lib/MockClient.ts
+++ b/src/lib/MockClient.ts
@@ -36,4 +36,18 @@ export class MockClient extends MockPeer {
       this.wsServer.emit('connection', this.peerWebSocket, incomingMessage);
     });
   }
+
+  /**
+   * Connect to server, run `callback` and close the connection.
+   *
+   * @param callback The function to execute whilst connected to server
+   */
+  public async use(callback: () => Promise<void>): Promise<void> {
+    await this.connect();
+    try {
+      await callback();
+    } finally {
+      this.close();
+    }
+  }
 }