feat: Implement mock server (#6)

README.md

@@ -13,33 +13,48 @@ npm install @relaycorp/ws-mock
 You should initialise `MockClient` by passing the `ws` server to be tested and then call `client.connect()` to initiate the connection. From that point you can interact with the server. For example:
 
 ```javascript
-  test('Challenge should be sent as soon as client connects', async () => {
-    const client = new MockClient(wsServer);
-    await client.connect();
+test('Challenge should be sent as soon as client connects', async () => {
+  const client = new MockClient(wsServer);
+  await client.connect();
 
-    const challengeSerialized = await client.receive();
-    expect(challengeSerialized).toBeInstanceOf(ArrayBuffer);
-  });
+  const challengeSerialized = await client.receive();
+  expect(challengeSerialized).toBeInstanceOf(ArrayBuffer);
+});
 ```
 
 You'll find [real-world examples in relaycorp/relaynet-internet-gateway](https://github.com/relaycorp/relaynet-internet-gateway/search?l=TypeScript&q=%22%40relaycorp%2Fws-mock%22).
 
 ## Using the mock server
 
-Simply initialise `MockServer` with the `ws` client connection you wish to test. For example:
+You basically need to initialise `MockServer` and replace the default export from `ws` with a mock WebSocket. Here's an example with Jest:
 
 ```javascript
-  test('Server message should be played back', async () => {
-    const mockConnection = new MockConnection();
-    const mockServer = new MockServer(mockConnection);
-    const clientUnderTest = new ClientUnderTest(mockConnection);
-
-    clientUnderTest.connectToServer();
-    mockServer.send('foo');
-
-    const clientResponse = await mockServer.receive();
-    expect(clientResponse).toEqual('foo');
-  });
+let mockServer: MockServer;
+beforeEach(() => {
+  mockServer = new MockServer();
+});
+jest.mock('ws', () => ({
+  __esModule: true,
+  default: jest.fn().mockImplementation(() => mockServer.mockClientWebSocket),
+}));
+
+test('Server message should be played back', async () => {
+  const clientUnderTest = new ClientUnderTest();
+  const messageToEcho = 'foo';
+
+  await Promise.all([
+    clientUnderTest.connectToServerAndInteractWithIt(),
+    // Configure the mock server to accept the incoming connection and return a message straightaway
+    mockServer.runActions(
+      new AcceptConnectionAction(),
+      new SendMessageAction(messageToEcho),
+    ),
+  ]);
+
+  // Check that the client sent the message back to the server:
+  const clientResponse = await mockServer.receive();
+  expect(clientResponse).toEqual(messageToEcho);
+});
 ```
 
 You'll find [real-world examples in relaycorp/relaynet-poweb-js](https://github.com/relaycorp/relaynet-poweb-js/search?l=TypeScript&q=%22%40relaycorp%2Fws-mock%22).
@@ -52,5 +67,7 @@ When using streams in the unit under test, make sure to mock the `createWebSocke
 import { createMockWebSocketStream } from '@relaycorp/ws-mock';
 import WebSocket from 'ws';
 
-jest.spyOn(WebSocket, 'createWebSocketStream').mockImplementation(createMockWebSocketStream);
+jest
+  .spyOn(WebSocket, 'createWebSocketStream')
+  .mockImplementation(createMockWebSocketStream);
 ```

src/index.ts

@@ -1,4 +1,6 @@
 export { MockClient } from './lib/MockClient';
 export { MockPeer } from './lib/MockPeer';
+export { MockServer } from './lib/MockServer';
+export * from './lib/MockServerAction';
 export { CloseFrame } from './lib/CloseFrame';
 export { createMockWebSocketStream } from './lib/stream';

src/lib/MockPeer.ts

@@ -1,30 +1,31 @@
-import { EventEmitter } from 'events';
-
 import { Data as WSData } from 'ws';
+
 import { CloseFrame } from './CloseFrame';
 import { MockWebSocket } from './MockWebSocket';
 
-export abstract class MockPeer extends EventEmitter {
+export abstract class MockPeer {
   protected readonly peerWebSocket = new MockWebSocket();
 
   get wasConnectionClosed(): boolean {
     return this.peerWebSocket.closeFrame !== null;
   }
 
-  public disconnect(code?: number, reason?: string): void {
+  public close(code?: number, reason?: string): void {
     this.peerWebSocket.emit('close', code, reason);
   }
 
   public abort(error: Error): void {
     this.peerWebSocket.emit('error', error);
   }
 
-  public async send(message: Buffer | string): Promise<void> {
+  public async send(message: WSData): Promise<void> {
     this.requireConnectionStillOpen();
 
+    const messageSerialized =
+      typeof message === 'string' ? message : this.convertBinaryType(message);
     return new Promise((resolve) => {
       this.peerWebSocket.once('message', resolve);
-      this.peerWebSocket.emit('message', message);
+      this.peerWebSocket.emit('message', messageSerialized);
     });
   }
 
@@ -47,6 +48,21 @@ export abstract class MockPeer extends EventEmitter {
     return this.peerWebSocket.getCloseFrameWhenAvailable();
   }
 
+  get peerCloseFrame(): CloseFrame | null {
+    return this.peerWebSocket.closeFrame;
+  }
+
+  protected convertBinaryType(
+    message: Buffer | ArrayBuffer | readonly Buffer[],
+  ): Buffer | ArrayBuffer | readonly Buffer[] {
+    const binaryType = this.peerWebSocket.binaryType;
+    if (binaryType === 'nodebuffer') {
+      return Buffer.isBuffer(message) ? message : Buffer.from(message);
+    }
+
+    throw new Error(`Unsupported WebSocket.binaryType (${binaryType}); feel free to open a PR`);
+  }
+
   private requireConnectionStillOpen(): void {
     if (this.wasConnectionClosed) {
       throw new Error('Connection was already closed');

src/lib/MockServer.ts

@@ -1,3 +1,22 @@
+import WebSocket from 'ws';
 import { MockPeer } from './MockPeer';
+import { MockServerAction } from './MockServerAction';
 
-export class MockServer extends MockPeer {}
+export class MockServer extends MockPeer {
+  get mockClientWebSocket(): WebSocket {
+    return this.peerWebSocket as any;
+  }
+
+  public async runActions(...actions: readonly MockServerAction[]): Promise<void> {
+    for (const action of actions) {
+      await action.run(this);
+    }
+  }
+
+  public async acceptConnection(): Promise<void> {
+    await new Promise((resolve) => {
+      this.peerWebSocket.once('open', resolve);
+      this.peerWebSocket.emit('open');
+    });
+  }
+}

src/lib/MockServerAction.ts

@@ -0,0 +1,48 @@
+// tslint:disable:max-classes-per-file
+
+import { Data } from 'ws';
+import { CloseFrame } from './CloseFrame';
+import { MockServer } from './MockServer';
+
+export abstract class MockServerAction {
+  // tslint:disable-next-line:readonly-keyword
+  protected _wasRun = false;
+
+  get wasRun(): boolean {
+    return this._wasRun;
+  }
+
+  public async run(_mockServer: MockServer): Promise<void> {
+    // tslint:disable-next-line:no-object-mutation
+    this._wasRun = true;
+  }
+}
+
+export class AcceptConnectionAction extends MockServerAction {
+  public async run(mockServer: MockServer): Promise<void> {
+    await mockServer.acceptConnection();
+    await super.run(mockServer);
+  }
+}
+
+export class CloseConnectionAction extends MockServerAction {
+  constructor(protected readonly closeFrame?: CloseFrame) {
+    super();
+  }
+
+  public async run(mockServer: MockServer): Promise<void> {
+    await mockServer.close(this.closeFrame?.code, this.closeFrame?.reason);
+    await super.run(mockServer);
+  }
+}
+
+export class SendMessageAction extends MockServerAction {
+  constructor(protected readonly message: Data) {
+    super();
+  }
+
+  public async run(mockServer: MockServer): Promise<void> {
+    await mockServer.send(this.message);
+    await super.run(mockServer);
+  }
+}

src/lib/MockWebSocket.ts

@@ -5,6 +5,9 @@ import { Data as WSData } from 'ws';
 import { CloseFrame } from './CloseFrame';
 
 export class MockWebSocket extends EventEmitter {
+  // tslint:disable-next-line:readonly-keyword
+  public binaryType: 'nodebuffer' | 'arraybuffer' = 'nodebuffer';
+
   // tslint:disable-next-line:readonly-keyword
   protected ownCloseFrame: CloseFrame | null = null;
   // tslint:disable-next-line:readonly-array