test: split Solana WebSocket Mock setup logic from Websocket Mocks cp-13.3.1

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;
+        }
       }
     });
   });