test: split Solana WebSocket Mock setup logic from Websocket Mocks (#35552)

<!--
Please submit this PR as a draft initially.
Do not mark it as "Ready for review" until the template has been
completely filled out, and PR status checks have passed at least once.
-->

## **Description**

As a follow-up of [this
PR](#35425), which
separated the logic for setting up the Websocket, from the test steps,
now this PR separates the logic for handling mocked wss requests from
the setup websocket server.

This means that now we can pass an array of websocket mocks.


Note, that the websocket mocks are different from the regular Http
mocks, they look like this:

```
export interface WebSocketMessageMock {
  messageIncludes: string | string[];
  response: object;
  delay?: number;
  logMessage?: string;
}
```

and the logic for handling the mocks is also different, as we are
handling this in our local websocket server. Meaning:

1. From Mocktt default mocks --> we check if we need a Websocket Solana
server, if so, we re-route any wss to the local wss
2. We setup the local Websocket server
3. We add 'onMessage' listeners to intercept the requests we want to
mock, and we send back the desired response



[![Open in GitHub
Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/MetaMask/metamask-extension/pull/35552?quickstart=1)

## **Changelog**

<!--
If this PR is not End-User-Facing and should not show up in the
CHANGELOG, you can choose to either:
1. Write `CHANGELOG entry: null`
4. Label with `no-changelog`

If this PR is End-User-Facing, please write a short User-Facing
description in the past tense like:
`CHANGELOG entry: Added a new tab for users to see their NFTs`
`CHANGELOG entry: Fixed a bug that was causing some NFTs to flicker`

(This helps the Release Engineer do their job more quickly and
accurately)
-->

CHANGELOG entry:

## **Related issues**

Fixes:

## **Manual testing steps**

1. Go to this page...
2.
5.

## **Screenshots/Recordings**

<!-- If applicable, add screenshots and/or recordings to visualize the
before and after of your change. -->

### **Before**

<!-- [screenshots/recordings] -->

### **After**

<!-- [screenshots/recordings] -->

## **Pre-merge author checklist**

- [ ] I've followed [MetaMask Contributor
Docs](https://github.com/MetaMask/contributor-docs) and [MetaMask
Extension Coding
Standards](https://github.com/MetaMask/metamask-extension/blob/main/.github/guidelines/CODING_GUIDELINES.md).
- [x] I've completed the PR template to the best of my ability
- [x] I’ve included tests if applicable
- [x] I’ve documented my code using [JSDoc](https://jsdoc.app/) format
if applicable
- [x] I’ve applied the right labels on the PR (see [labeling
guidelines](https://github.com/MetaMask/metamask-extension/blob/main/.github/guidelines/LABELING_GUIDELINES.md)).
Not required for external contributors.

## **Pre-merge reviewer checklist**

- [ ] I've manually tested the PR (e.g. pull and build branch, run the
app, test code being changed).
- [ ] I confirm that this PR addresses all acceptance criteria described
in the ticket it closes and includes the necessary testing evidence such
as recordings and or screenshots.

Author: seaona

test/e2e/helpers.js

@@ -26,7 +26,7 @@ const {
   getServerMochaToBackground,
 } = require('./background-socket/server-mocha-to-background');
 const LocalWebSocketServer = require('./websocket-server').default;
-const { setSolanaWebsocketMocks } = require('./websocket-solana-mocks');
+const { setupSolanaWebsocketMocks } = require('./websocket-solana-mocks');
 
 const tinyDelayMs = 200;
 const regularDelayMs = tinyDelayMs * 2;
@@ -128,7 +128,10 @@ async function withFixtures(options, testSuite) {
     ethConversionInUsd,
     monConversionInUsd,
     manifestFlags,
-    withSolanaWebSocket = false,
+    withSolanaWebSocket = {
+      server: false,
+      mocks: [],
+    },
   } = options;
 
   // Normalize localNodeOptions
@@ -261,12 +264,10 @@ async function withFixtures(options, testSuite) {
       }
     }
 
-    if (withSolanaWebSocket) {
+    if (withSolanaWebSocket.server) {
       localWebSocketServer = LocalWebSocketServer.getServerInstance();
       localWebSocketServer.start();
-      // All specs use the same ws mocks.
-      // If we need custom ws mocks we can expand logic for supporting custom ws mocks like with http
-      await setSolanaWebsocketMocks();
+      await setupSolanaWebsocketMocks(withSolanaWebSocket.mocks);
     }
 
     const { mockedEndpoint, getPrivacyReport } = await setupMocking(
@@ -454,7 +455,7 @@ async function withFixtures(options, testSuite) {
         })(),
       );
 
-      if (withSolanaWebSocket) {
+      if (withSolanaWebSocket.server) {
         shutdownTasks.push(localWebSocketServer.stopAndCleanup());
       }
 

test/e2e/mock-e2e.js

@@ -143,7 +143,7 @@ const privateHostMatchers = [
  * @param {object} options - Network mock options.
  * @param {string} options.chainId - The chain ID used by the default configured network.
  * @param {string} options.ethConversionInUsd - The USD conversion rate for ETH.
- * @param {boolean} withSolanaWebSocket - If we want to re-route all the ws requests to our Solana Local WS server
+ * @param {object} withSolanaWebSocket - Solana WebSocket configuration with server flag and mocks function
  * @returns {Promise<SetupMockReturn>}
  */
 async function setupMocking(
@@ -946,7 +946,7 @@ async function setupMocking(
    * Solana Websocket
    * Setup HTTP intercept for WebSocket handshake requests
    */
-  if (withSolanaWebSocket) {
+  if (withSolanaWebSocket.server) {
     await server
       .forAnyWebSocket()
       .matching((req) =>

test/e2e/tests/solana/common-solana.ts

@@ -11,6 +11,7 @@ import { ACCOUNT_TYPE } from '../../constants';
 import { loginWithBalanceValidation } from '../../page-objects/flows/login.flow';
 import { mockProtocolSnap } from '../../mock-response-data/snaps/snap-binary-mocks';
 import AssetListPage from '../../page-objects/pages/home/asset-list';
+import { DEFAULT_SOLANA_WS_MOCKS } from './mocks/websocketDefaultMocks';
 
 const SOLANA_URL_REGEX_MAINNET =
   /^https:\/\/solana-(mainnet|devnet)\.infura\.io\/v3*/u;
@@ -1611,7 +1612,10 @@ export async function withSolanaAccountSnap(
       fixtures: fixtures.build(),
       title,
       dapp: true,
-      withSolanaWebSocket: true,
+      withSolanaWebSocket: {
+        server: true,
+        mocks: DEFAULT_SOLANA_WS_MOCKS,
+      },
       manifestFlags: {
         // This flag is used to enable/disable the remote mode for the carousel
         // component, which will impact to the slides count.

test/e2e/tests/solana/mocks/websocketDefaultMocks.ts

@@ -0,0 +1,65 @@
+/**
+ * Configuration for a WebSocket message mock
+ */
+export type WebSocketMessageMock = {
+  /** String(s) that the message should include to trigger this mock */
+  messageIncludes: string | string[];
+  /** The JSON response to send back */
+  response: object;
+  /** Delay before sending the response (in milliseconds) */
+  delay?: number;
+  /** Custom log message for this mock */
+  logMessage?: string;
+};
+
+export const DEFAULT_SOLANA_WS_MOCKS: WebSocketMessageMock[] = [
+  {
+    messageIncludes: 'signatureSubscribe',
+    response: {
+      jsonrpc: '2.0',
+      result: 8648699534240963,
+      id: '1',
+    },
+    delay: 500,
+    logMessage: 'Signature subscribe message received from client',
+  },
+  {
+    messageIncludes: 'accountSubscribe',
+    response: {
+      jsonrpc: '2.0',
+      result:
+        'b07ebf7caf2238a9b604d4dfcaf1934280fcd347d6eded62bc0def6cbb767d11',
+      id: '1',
+    },
+    delay: 500,
+    logMessage: 'Account subscribe message received from client',
+  },
+  {
+    messageIncludes: [
+      'programSubscribe',
+      'TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA',
+    ],
+    response: {
+      jsonrpc: '2.0',
+      result:
+        '568eafd45635c108d0d426361143de125a841628a58679f5a024cbab9a20b41c',
+      id: '1',
+    },
+    delay: 500,
+    logMessage: 'Program subscribe message received from client',
+  },
+  {
+    messageIncludes: [
+      'programSubscribe',
+      'TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb',
+    ],
+    response: {
+      jsonrpc: '2.0',
+      result:
+        'f33dd9975158af47bf16c7f6062a73191d4595c59cfec605d5a51e25c65ffb51',
+      id: '1',
+    },
+    delay: 500,
+    logMessage: 'Program subscribe message received from client',
+  },
+];

test/e2e/tests/solana/web-socket-connection.spec.ts

@@ -8,14 +8,18 @@ import HeaderNavbar from '../../page-objects/pages/header-navbar';
 import AccountListPage from '../../page-objects/pages/account-list-page';
 import FixtureBuilder from '../../fixture-builder';
 import LocalWebSocketServer from '../../websocket-server';
+import { DEFAULT_SOLANA_WS_MOCKS } from './mocks/websocketDefaultMocks';
 
 describe('Solana Web Socket', function (this: Suite) {
   it('a websocket connection is open when MetaMask full view is open', async function () {
     await withFixtures(
       {
         fixtures: new FixtureBuilder().build(),
         title: this.test?.fullTitle(),
-        withSolanaWebSocket: true,
+        withSolanaWebSocket: {
+          server: true,
+          mocks: DEFAULT_SOLANA_WS_MOCKS,
+        },
         manifestFlags: {
           remoteFeatureFlags: {
             addSolanaAccount: true,
@@ -50,7 +54,10 @@ describe('Solana Web Socket', function (this: Suite) {
       {
         fixtures: new FixtureBuilder().build(),
         title: this.test?.fullTitle(),
-        withSolanaWebSocket: true,
+        withSolanaWebSocket: {
+          server: true,
+          mocks: DEFAULT_SOLANA_WS_MOCKS,
+        },
         manifestFlags: {
           remoteFeatureFlags: {
             addSolanaAccount: true,
@@ -95,7 +102,10 @@ describe('Solana Web Socket', function (this: Suite) {
       {
         fixtures: new FixtureBuilder().build(),
         title: this.test?.fullTitle(),
-        withSolanaWebSocket: true,
+        withSolanaWebSocket: {
+          server: true,
+          mocks: DEFAULT_SOLANA_WS_MOCKS,
+        },
         manifestFlags: {
           remoteFeatureFlags: {
             addSolanaAccount: true,

test/e2e/websocket-solana-mocks.ts

@@ -1,12 +1,16 @@
 // eslint-disable-next-line @typescript-eslint/no-shadow
 import { WebSocket } from 'ws';
 import LocalWebSocketServer from './websocket-server';
+import { WebSocketMessageMock } from './tests/solana/mocks/websocketDefaultMocks';
 
 /**
- * WebSocket Solana mocks
- * This function should be called after the WebSocket server is started
+ * Sets up Solana WebSocket mocks with configurable message handlers
+ *
+ * @param mocks - Array of message mock configurations
  */
-export async function setSolanaWebsocketMocks(): Promise<void> {
+export async function setupSolanaWebsocketMocks(
+  mocks: WebSocketMessageMock[] = [],
+): Promise<void> {
   const localWebSocketServer = LocalWebSocketServer.getServerInstance();
   const wsServer = localWebSocketServer.getServer();
 
@@ -18,72 +22,34 @@ export async function setSolanaWebsocketMocks(): Promise<void> {
     socket.on('message', (data) => {
       const message = data.toString();
       console.log('Message received from client:', message);
-      if (message.includes('signatureSubscribe')) {
-        console.log('Signature subscribe message received from client');
-        setTimeout(() => {
-          socket.send(
-            JSON.stringify({
-              jsonrpc: '2.0',
-              result: 8648699534240963,
-              id: '1',
-            }),
-          );
-          console.log('Simulated message sent to the client');
-        }, 500); // Delay the message by 500ms
-      }
-      if (message.includes('accountSubscribe')) {
-        console.log('Account subscribe message received from client');
-        setTimeout(() => {
-          socket.send(
-            JSON.stringify({
-              jsonrpc: '2.0',
-              result:
-                'b07ebf7caf2238a9b604d4dfcaf1934280fcd347d6eded62bc0def6cbb767d11',
-              id: '1',
-            }),
-          );
-          console.log(
-            'Simulated message for accountSubscribe sent to the client',
-          );
-        }, 500); // Delay the message by 500ms
-      }
-      if (
-        message.includes('programSubscribe') &&
-        message.includes('TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA')
-      ) {
-        console.log('Program subscribe message received from client');
-        setTimeout(() => {
-          socket.send(
-            JSON.stringify({
-              jsonrpc: '2.0',
-              result:
-                '568eafd45635c108d0d426361143de125a841628a58679f5a024cbab9a20b41c',
-              id: '1',
-            }),
-          );
-          console.log(
-            'Simulated message for programSubscribe Token2022 sent to the client',
-          );
-        }, 500); // Delay the message by 500ms
-      }
-      if (
-        message.includes('programSubscribe') &&
-        message.includes('TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb')
-      ) {
-        console.log('Program subscribe message received from client');
-        setTimeout(() => {
-          socket.send(
-            JSON.stringify({
-              jsonrpc: '2.0',
-              result:
-                'f33dd9975158af47bf16c7f6062a73191d4595c59cfec605d5a51e25c65ffb51',
-              id: '1',
-            }),
-          );
-          console.log(
-            'Simulated message for programSubscribe sent to the client',
-          );
-        }, 500); // Delay the message by 500ms
+
+      // Check each mock configuration
+      for (const mock of mocks) {
+        const includes = Array.isArray(mock.messageIncludes)
+          ? mock.messageIncludes
+          : [mock.messageIncludes];
+
+        // Check if all required strings are included in the message
+        const matches = includes.every((includeStr) =>
+          message.includes(includeStr),
+        );
+
+        if (matches) {
+          if (mock.logMessage) {
+            console.log(mock.logMessage);
+          }
+
+          const delay = mock.delay || 500;
+          setTimeout(() => {
+            socket.send(JSON.stringify(mock.response));
+            console.log(
+              `Simulated message sent to the client for: ${includes.join(' + ')}`,
+            );
+          }, delay);
+
+          // Break after first match to avoid multiple responses
+          break;
+        }
       }
     });
   });