|
| 1 | +# WebSocket Mocking in E2E Tests |
| 2 | + |
| 3 | +The E2E framework includes a local WebSocket server that intercepts production WebSocket connections during tests. This enables deterministic testing of features that use real-time data streams (e.g., account activity notifications). |
| 4 | + |
| 5 | +## Architecture |
| 6 | + |
| 7 | +``` |
| 8 | +App (E2E build) |
| 9 | + │ |
| 10 | + ├─ shim.js rewrites wss://gateway.api.cx.metamask.io → ws://localhost:<port> |
| 11 | + │ |
| 12 | + └─ LocalWebSocketServer (tests/websocket/server.ts) |
| 13 | + │ |
| 14 | + └─ Protocol mock (e.g. account-activity-mocks.ts) |
| 15 | + └─ Handles subscribe/unsubscribe, sends mock notifications |
| 16 | +``` |
| 17 | + |
| 18 | +**Key files:** |
| 19 | + |
| 20 | +| File | Purpose | |
| 21 | +| ------------------------------------------- | -------------------------------------------------------------------------------- | |
| 22 | +| `tests/websocket/server.ts` | `LocalWebSocketServer` — generic WS server implementing the `Resource` interface | |
| 23 | +| `tests/websocket/constants.ts` | Service definitions (production URL, fallback port, launch arg key) | |
| 24 | +| `tests/websocket/account-activity-mocks.ts` | Protocol-specific mock for AccountActivity WS | |
| 25 | +| `tests/framework/fixtures/FixtureHelper.ts` | Creates and manages the WS server per test run | |
| 26 | + |
| 27 | +## How It Works |
| 28 | + |
| 29 | +1. `withFixtures` in `FixtureHelper.ts` creates a `LocalWebSocketServer` instance for every test run |
| 30 | +2. The server is started via `startResourceWithRetry` with automatic port allocation |
| 31 | +3. `setupAccountActivityMocks()` attaches protocol-specific message handlers |
| 32 | +4. The allocated port is passed to the app via launch args (`accountActivityWsPort`) |
| 33 | +5. The app's E2E shim rewrites the production WebSocket URL to `ws://localhost:<port>` |
| 34 | +6. After the test, the server is stopped and connections are cleaned up |
| 35 | + |
| 36 | +## Using the Existing AccountActivity Mock |
| 37 | + |
| 38 | +The AccountActivity WebSocket mock is available in every test automatically. For most tests, you don't need to do anything — the mock server handles subscribe/unsubscribe handshakes in the background. |
| 39 | + |
| 40 | +To explicitly test WebSocket behavior, use the helpers from `account-activity-mocks.ts`: |
| 41 | + |
| 42 | +```typescript |
| 43 | +import { |
| 44 | + waitForAccountActivitySubscription, |
| 45 | + getAccountActivitySubscriptionCount, |
| 46 | + waitForAccountActivityDisconnection, |
| 47 | + createBalanceUpdateNotification, |
| 48 | +} from '../../websocket/account-activity-mocks'; |
| 49 | +``` |
| 50 | + |
| 51 | +### Waiting for a subscription |
| 52 | + |
| 53 | +```typescript |
| 54 | +// Set up the waiter BEFORE the action that triggers the connection |
| 55 | +const subscriptionPromise = waitForAccountActivitySubscription(); |
| 56 | +await loginToApp(); |
| 57 | +const subscriptionId = await subscriptionPromise; |
| 58 | +``` |
| 59 | + |
| 60 | +### Checking subscription count |
| 61 | + |
| 62 | +```typescript |
| 63 | +assertEqual(getAccountActivitySubscriptionCount(), 1); |
| 64 | +``` |
| 65 | + |
| 66 | +### Waiting for disconnection |
| 67 | + |
| 68 | +```typescript |
| 69 | +await device.sendToHome(); |
| 70 | +await waitForAccountActivityDisconnection(); |
| 71 | +``` |
| 72 | + |
| 73 | +### Sending a mock balance update |
| 74 | + |
| 75 | +```typescript |
| 76 | +const notification = createBalanceUpdateNotification({ |
| 77 | + subscriptionId, |
| 78 | + channel: 'account-activity.v1', |
| 79 | + address: '0x1234...', |
| 80 | + chain: 'eip155:1', |
| 81 | + updates: [ |
| 82 | + { |
| 83 | + asset: { fungible: true, type: 'native', unit: 'ETH', decimals: 18 }, |
| 84 | + postBalance: { amount: '1000000000000000000' }, |
| 85 | + transfers: [ |
| 86 | + { from: '0x0...', to: '0x1234...', amount: '500000000000000000' }, |
| 87 | + ], |
| 88 | + }, |
| 89 | + ], |
| 90 | +}); |
| 91 | + |
| 92 | +// Broadcast to all connected clients |
| 93 | +accountActivityWsServer.sendMessage(JSON.stringify(notification)); |
| 94 | +``` |
| 95 | + |
| 96 | +### Example spec |
| 97 | + |
| 98 | +See `tests/smoke/account-activity/web-socket-connection.spec.ts` for a complete example with three tests covering subscribe-on-login, resubscribe-after-background, and resubscribe-after-lock. |
| 99 | + |
| 100 | +## Adding a New WebSocket Service Mock |
| 101 | + |
| 102 | +To mock a new WebSocket service (e.g., a new real-time data feed): |
| 103 | + |
| 104 | +### 1. Define the service in `tests/websocket/constants.ts` |
| 105 | + |
| 106 | +```typescript |
| 107 | +export const MY_NEW_WS: WebSocketServiceConfig = { |
| 108 | + url: 'wss://my-service.example.com', |
| 109 | + fallbackPort: 8090, // pick a unique port |
| 110 | + launchArgKey: 'myNewWsPort', |
| 111 | +}; |
| 112 | + |
| 113 | +// Add it to the services array |
| 114 | +export const WS_SERVICES: WebSocketServiceConfig[] = [ |
| 115 | + ACCOUNT_ACTIVITY_WS, |
| 116 | + MY_NEW_WS, |
| 117 | +]; |
| 118 | +``` |
| 119 | + |
| 120 | +### 2. Create a protocol mock in `tests/websocket/` |
| 121 | + |
| 122 | +```typescript |
| 123 | +// tests/websocket/my-service-mocks.ts |
| 124 | +import type LocalWebSocketServer from './server.ts'; |
| 125 | +import { WebSocket } from 'ws'; |
| 126 | + |
| 127 | +export async function setupMyServiceMocks( |
| 128 | + server: LocalWebSocketServer, |
| 129 | +): Promise<void> { |
| 130 | + const wsServer = server.getServer(); |
| 131 | + |
| 132 | + wsServer.on('connection', (socket: WebSocket) => { |
| 133 | + // Handle incoming messages |
| 134 | + socket.on('message', (data) => { |
| 135 | + const raw = data.toString(); |
| 136 | + // Parse and respond based on your protocol |
| 137 | + }); |
| 138 | + |
| 139 | + // Send initial data on connection |
| 140 | + socket.send(JSON.stringify({ type: 'connected' })); |
| 141 | + }); |
| 142 | +} |
| 143 | +``` |
| 144 | + |
| 145 | +### 3. Register in `FixtureHelper.ts` |
| 146 | + |
| 147 | +Add the server creation, startup, and cleanup alongside the existing `accountActivityWsServer` pattern: |
| 148 | + |
| 149 | +```typescript |
| 150 | +const myNewWsServer = new LocalWebSocketServer('myNewService'); |
| 151 | +// ... |
| 152 | +await startResourceWithRetry(ResourceType.MY_NEW_WS, myNewWsServer); |
| 153 | +await setupMyServiceMocks(myNewWsServer); |
| 154 | +``` |
| 155 | + |
| 156 | +Pass the port via launch args: |
| 157 | + |
| 158 | +```typescript |
| 159 | +[MY_NEW_WS.launchArgKey]: isAndroid |
| 160 | + ? `${MY_NEW_WS.fallbackPort}` |
| 161 | + : `${myNewWsServer.getServerPort()}`, |
| 162 | +``` |
| 163 | + |
| 164 | +### 4. Register the port in PortManager |
| 165 | + |
| 166 | +Add the new `ResourceType` and fallback port in `tests/framework/PortManager.ts` so the framework can manage port allocation and Android `adb reverse` forwarding. |
0 commit comments