feat(playwright): grab WS messages (#3789)

Author: kobenguyent

docs/helpers/Playwright.md

@@ -861,6 +861,10 @@ I.fillField({css: 'form#login input[name=username]'}, 'John');
 
 Resets all recorded network requests.
 
+### flushWebSocketMessages
+
+Resets all recorded WS messages.
+
 ### focus
 
 Calls [focus][10] on the matching element.
@@ -1257,6 +1261,12 @@ let inputs = await I.grabValueFromAll('//form/input');
 
 Returns **[Promise][20]<[Array][22]<[string][8]>>** attribute value
 
+### grabWebSocketMessages
+
+Grab the recording WS messages
+
+Returns **[Array][22]<any>** 
+
 ### handleDownloads
 
 Handles a file download. A file name is required to save the file on disk.
@@ -1960,6 +1970,15 @@ I.startRecordingTraffic();
 
 Returns **[Promise][20]<void>** 
 
+### startRecordingWebSocketMessages
+
+Starts recording of websocket messages.
+This also resets recorded websocket messages.
+
+```js
+await I.startRecordingWebSocketMessages();
+```
+
 ### stopMockingRoute
 
 Stops network mocking created by `mockRoute`.
@@ -1984,6 +2003,14 @@ Stops recording of network traffic. Recorded traffic is not flashed.
 I.stopRecordingTraffic();
 ```
 
+### stopRecordingWebSocketMessages
+
+Stops recording WS messages. Recorded WS messages is not flashed.
+
+```js
+await I.stopRecordingWebSocketMessages();
+```
+
 ### switchTo
 
 Switches frame or in case of null locator reverts to parent.

lib/helper/Playwright.js

@@ -318,6 +318,12 @@ class Playwright extends Helper {
     this.recording = false;
     this.recordedAtLeastOnce = false;
 
+    // for websocket messages
+    this.webSocketMessages = [];
+    this.recordingWebSocketMessages = false;
+    this.recordedWebSocketMessagesAtLeastOnce = false;
+    this.cdpSession = null;
+
     // override defaults with config
     this._setConfig(config);
   }
@@ -2981,6 +2987,102 @@ class Playwright extends Helper {
     });
     return dumpedTraffic;
   }
+
+  /**
+   * Starts recording of websocket messages.
+   * This also resets recorded websocket messages.
+   *
+   * ```js
+   * await I.startRecordingWebSocketMessages();
+   * ```
+   *
+   */
+  async startRecordingWebSocketMessages() {
+    this.flushWebSocketMessages();
+    this.recordingWebSocketMessages = true;
+    this.recordedWebSocketMessagesAtLeastOnce = true;
+
+    this.cdpSession = await this.getNewCDPSession();
+    await this.cdpSession.send('Network.enable');
+    await this.cdpSession.send('Page.enable');
+
+    this.cdpSession.on(
+      'Network.webSocketFrameReceived',
+      (payload) => {
+        this._logWebsocketMessages(this._getWebSocketLog('RECEIVED', payload));
+      },
+    );
+
+    this.cdpSession.on(
+      'Network.webSocketFrameSent',
+      (payload) => {
+        this._logWebsocketMessages(this._getWebSocketLog('SENT', payload));
+      },
+    );
+
+    this.cdpSession.on(
+      'Network.webSocketFrameError',
+      (payload) => {
+        this._logWebsocketMessages(this._getWebSocketLog('ERROR', payload));
+      },
+    );
+  }
+
+  /**
+   * Stops recording WS messages. Recorded WS messages is not flashed.
+   *
+   * ```js
+   * await I.stopRecordingWebSocketMessages();
+   * ```
+   */
+  async stopRecordingWebSocketMessages() {
+    await this.cdpSession.send('Network.disable');
+    await this.cdpSession.send('Page.disable');
+    this.page.removeAllListeners('Network');
+    this.recordingWebSocketMessages = false;
+  }
+
+  /**
+   *  Grab the recording WS messages
+   *
+   * @return { Array<any> }
+   *
+   */
+  grabWebSocketMessages() {
+    if (!this.recordingWebSocketMessages) {
+      if (!this.recordedWebSocketMessagesAtLeastOnce) {
+        throw new Error('Failure in test automation. You use "I.grabWebSocketMessages", but "I.startRecordingWebSocketMessages" was never called before.');
+      }
+    }
+    return this.webSocketMessages;
+  }
+
+  /**
+   * Resets all recorded WS messages.
+   */
+  flushWebSocketMessages() {
+    this.webSocketMessages = [];
+  }
+
+  _getWebSocketMessage(payload) {
+    if (payload.errorMessage) {
+      return payload.errorMessage;
+    }
+
+    return payload.response.payloadData;
+  }
+
+  _getWebSocketLog(prefix, payload) {
+    return `${prefix} ID: ${payload.requestId} TIMESTAMP: ${payload.timestamp} (${new Date().toISOString()})\n\n${this._getWebSocketMessage(payload)}\n\n`;
+  }
+
+  async getNewCDPSession() {
+    return this.page.context().newCDPSession(this.page);
+  }
+
+  _logWebsocketMessages(message) {
+    this.webSocketMessages += message;
+  }
 }
 
 module.exports = Playwright;

test/helper/Playwright_test.js

@@ -938,6 +938,47 @@ describe('Playwright', function () {
     });
   });
 
+  describe('#startRecordingWebSocketMessages, #grabWebSocketMessages, #stopRecordingWebSocketMessages', () => {
+    it('should throw error when calling grabWebSocketMessages before startRecordingWebSocketMessages', () => {
+      try {
+        I.amOnPage('https://websocketstest.com/');
+        I.waitForText('Work for You!');
+        I.grabWebSocketMessages();
+      } catch (e) {
+        expect(e.message).to.equal('Failure in test automation. You use "I.grabWebSocketMessages", but "I.startRecordingWebSocketMessages" was never called before.');
+      }
+    });
+
+    it('should flush the WS messages', async () => {
+      await I.startRecordingWebSocketMessages();
+      I.amOnPage('https://websocketstest.com/');
+      I.waitForText('Work for You!');
+      I.flushNetworkTraffics();
+      const wsMessages = I.grabWebSocketMessages();
+      expect(wsMessages.length).to.equal(0);
+    });
+
+    it('should see recording WS messages', async () => {
+      await I.startRecordingWebSocketMessages();
+      await I.amOnPage('https://websocketstest.com/');
+      I.waitForText('Work for You!');
+      const wsMessages = I.grabWebSocketMessages();
+      expect(wsMessages.length).to.greaterThan(0);
+    });
+
+    it('should not see recording WS messages', async () => {
+      await I.startRecordingWebSocketMessages();
+      await I.amOnPage('https://websocketstest.com/');
+      I.waitForText('Work for You!');
+      const wsMessages = I.grabWebSocketMessages();
+      await I.stopRecordingWebSocketMessages();
+      await I.amOnPage('https://websocketstest.com/');
+      I.waitForText('Work for You!');
+      const afterWsMessages = I.grabWebSocketMessages();
+      expect(wsMessages.length).to.equal(afterWsMessages.length);
+    });
+  });
+
   describe('#makeApiRequest', () => {
     it('should make 3rd party API request', async () => {
       const response = await I.makeApiRequest('get', 'https://reqres.in/api/users?page=2');