test: new switchToWindowWithTitle w/ Extension communication

app/scripts/background.js

@@ -43,6 +43,7 @@ import { checkForLastErrorAndLog } from '../../shared/modules/browser-runtime.ut
 import { isManifestV3 } from '../../shared/modules/mv3.utils';
 import { maskObject } from '../../shared/modules/object.utils';
 import { FIXTURE_STATE_METADATA_VERSION } from '../../test/e2e/default-fixture';
+import { getSocketBackgroundToMocha } from '../../test/e2e/background-socket/socket-background-to-mocha';
 import migrations from './migrations';
 import Migrator from './lib/migrator';
 import ExtensionPlatform from './platforms/extension';
@@ -280,6 +281,11 @@ async function initialize() {
 
     let isFirstMetaMaskControllerSetup;
 
+    // We only want to start this if we are running a test build, not for the release build
+    if (process.env.IN_TEST) {
+      getSocketBackgroundToMocha();
+    }
+
     if (isManifestV3) {
       // Save the timestamp immediately and then every `SAVE_TIMESTAMP_INTERVAL`
       // miliseconds. This keeps the service worker alive.

development/build/manifest.js

@@ -67,6 +67,7 @@ function createManifestTasks({
       ...manifest.permissions,
       'webRequestBlocking',
       'http://localhost/*',
+      'tabs', // test builds need tabs permission for switchToWindowWithTitle
     ];
   });
 
@@ -76,6 +77,7 @@ function createManifestTasks({
       ...manifest.permissions,
       'webRequestBlocking',
       'http://localhost/*',
+      'tabs', // test builds need tabs permission for switchToWindowWithTitle
     ];
   });
 

package.json

@@ -19,8 +19,8 @@
     "dist:mmi:debug": "yarn dist --build-type mmi --apply-lavamoat=false --snow=false",
     "build": "yarn lavamoat:build",
     "build:dev": "node development/build/index.js",
-    "start:test": "SEGMENT_HOST='https://api.segment.io' SEGMENT_WRITE_KEY='FAKE' SENTRY_DSN_DEV=https://[email protected]/0000000 yarn build:dev testDev",
-    "start:test:flask": "SEGMENT_HOST='https://api.segment.io' SEGMENT_WRITE_KEY='FAKE' SENTRY_DSN_DEV=https://[email protected]/0000000 yarn build:dev testDev --build-type flask --apply-lavamoat=false --snow=false",
+    "start:test": "SEGMENT_HOST='https://api.segment.io' SEGMENT_WRITE_KEY='FAKE' SENTRY_DSN_DEV=https://[email protected]/0000000 yarn build:dev testDev --apply-lavamoat=false",
+    "start:test:flask": "SEGMENT_HOST='https://api.segment.io' SEGMENT_WRITE_KEY='FAKE' SENTRY_DSN_DEV=https://[email protected]/0000000 yarn build:dev testDev --build-type flask --apply-lavamoat=false",
     "start:test:mv2:flask": "ENABLE_MV3=false SEGMENT_HOST='https://api.segment.io' SEGMENT_WRITE_KEY='FAKE' SENTRY_DSN_DEV=https://[email protected]/0000000 yarn build:dev testDev --build-type flask --apply-lavamoat=false --snow=false",
     "start:test:mv2": "ENABLE_MV3=false BLOCKAID_FILE_CDN=static.cx.metamask.io/api/v1/confirmations/ppom SEGMENT_HOST='https://api.segment.io' SEGMENT_WRITE_KEY='FAKE' SENTRY_DSN_DEV=https://[email protected]/0000000 yarn build:dev testDev --apply-lavamoat=false --snow=false",
     "benchmark:chrome": "SELENIUM_BROWSER=chrome ts-node test/e2e/benchmark.js",
@@ -489,6 +489,7 @@
     "@tsconfig/node20": "^20.1.2",
     "@types/babelify": "^7.3.7",
     "@types/browserify": "^12.0.37",
+    "@types/chrome": "^0.0.268",
     "@types/currency-formatter": "^1.5.1",
     "@types/fs-extra": "^9.0.13",
     "@types/gulp": "^4.0.9",
@@ -517,6 +518,7 @@
     "@types/w3c-web-hid": "^1.0.3",
     "@types/watchify": "^3.11.1",
     "@types/webextension-polyfill": "^0.10.4",
+    "@types/ws": "^8.5.10",
     "@types/yargs": "^17.0.32",
     "@typescript-eslint/eslint-plugin": "^7.10.0",
     "@typescript-eslint/parser": "^7.10.0",
@@ -640,6 +642,7 @@
     "watchify": "^4.0.0",
     "webextension-polyfill": "^0.8.0",
     "webpack": "^5.91.0",
+    "ws": "^8.17.1",
     "yaml": "^2.4.1",
     "yargs": "^17.7.2"
   },
@@ -707,6 +710,8 @@
       "@trezor/connect-web>@trezor/connect>@trezor/utxo-lib>tiny-secp256k1": false,
       "@storybook/test-runner>@swc/core": false,
       "@lavamoat/lavadome-react>@lavamoat/preinstall-always-fail": false,
+      "ws>bufferutil": false,
+      "ws>utf-8-validate": false,
       "tsx>esbuild": false,
       "@metamask/eth-trezor-keyring>@trezor/connect-web>@trezor/connect>@trezor/protobuf>protobufjs": false,
       "firebase>@firebase/firestore>@grpc/proto-loader>protobufjs": false,

test/e2e/accounts/common.ts

@@ -10,7 +10,6 @@ import {
   validateContractDetails,
   multipleGanacheOptions,
   regularDelayMs,
-  openDapp,
 } from '../helpers';
 import { Driver } from '../webdriver/driver';
 import { TEST_SNAPS_SIMPLE_KEYRING_WEBSITE_URL } from '../constants';
@@ -184,20 +183,8 @@ async function switchToAccount2(driver: Driver) {
 }
 
 export async function connectAccountToTestDapp(driver: Driver) {
-  try {
-    // Do an unusually fast switchToWindowWithTitle, just 1 second
-    await driver.switchToWindowWithTitle(
-      WINDOW_TITLES.TestDApp,
-      null,
-      1000,
-      1000,
-    );
-  } catch {
-    await driver.switchToWindowWithTitle(
-      WINDOW_TITLES.ExtensionInFullScreenView,
-    );
-    await openDapp(driver);
-  }
+  switchToOrOpenDapp(driver);
+
   await driver.clickElement('#connectButton');
 
   await driver.delay(regularDelayMs);

test/e2e/accounts/snap-account-signatures.spec.ts

@@ -1,5 +1,5 @@
 import { Suite } from 'mocha';
-import { openDapp, withFixtures } from '../helpers';
+import { withFixtures } from '../helpers';
 import { Driver } from '../webdriver/driver';
 import {
   accountSnapFixtures,
@@ -27,8 +27,6 @@ describe('Snap Account Signatures', function (this: Suite) {
 
           const newPublicKey = await makeNewAccountAndSwitch(driver);
 
-          await openDapp(driver);
-
           // Run all 6 signature types
           const locatorIDs = [
             '#ethSign',

test/e2e/background-socket/server-mocha-to-background.ts

@@ -0,0 +1,112 @@
+import events from 'events';
+import log from 'loglevel';
+import { WebSocketServer } from 'ws';
+import { WindowProperties } from './window-handles';
+
+/**
+ * This singleton class runs on the Mocha/Selenium test.
+ * It's used to communicate from the Mocha/Selenium test to the Extension background script (service worker in MV3).
+ */
+class ServerMochaToBackground {
+  private server: WebSocketServer;
+
+  private ws: WebSocket | null = null;
+
+  private eventEmitter;
+
+  constructor() {
+    this.server = new WebSocketServer({ port: 8111 });
+
+    log.debug('ServerMochaToBackground created');
+
+    this.server.on('connection', (ws: WebSocket) => {
+      this.ws = ws;
+
+      ws.onmessage = (ev: MessageEvent) => this.receivedMessage(ev.data);
+    });
+
+    this.eventEmitter = new events.EventEmitter();
+  }
+
+  // This function is never explicitly called, but in teh future it could be
+  stop() {
+    if (this.ws) {
+      this.ws.close();
+    }
+
+    this.server.close();
+
+    log.debug('ServerMochaToBackground stopped');
+  }
+
+  // Send a message to the Extension background script (service worker in MV3)
+  send(message: string | object) {
+    if (!this.ws) {
+      log.debug('No client connected');
+      return;
+    }
+
+    if (typeof message === 'string') {
+      this.ws.send(message);
+    } else {
+      this.ws.send(JSON.stringify(message));
+    }
+  }
+
+  // Handle messages received from the Extension background script (service worker in MV3)
+  private receivedMessage(message: string) {
+    let msg;
+
+    try {
+      msg = JSON.parse(message);
+    } catch (e) {
+      log.error('error in JSON', e);
+      return;
+    }
+
+    if (msg.command === 'openTabs') {
+      this.eventEmitter.emit('openTabs', msg);
+    } else if (msg.command === 'notFound') {
+      throw new Error('No window found by background script');
+    }
+  }
+
+  // This is not used in the current code, but could be used in the future
+  queryTabs(tabTitle: string) {
+    this.send({ command: 'queryTabs', title: tabTitle });
+  }
+
+  // Sends the message to the Extension, and waits for a response
+  async waitUntilWindowWithProperty(property: WindowProperties, value: string) {
+    this.send({ command: 'waitUntilWindowWithProperty', property, value });
+
+    const tabs = await this.waitForResponse();
+    log.debug('got the response', tabs);
+
+    // The return value here is less useful than we had hoped, because the tabs
+    // are not in the same order as driver.getAllWindowHandles()
+    return tabs;
+  }
+
+  // This is a way to wait for an event async, without timeouts or polling
+  async waitForResponse() {
+    return new Promise((resolve) => {
+      this.eventEmitter.on('openTabs', resolve);
+    });
+  }
+}
+
+// Singleton setup below
+let _serverMochaToBackground: ServerMochaToBackground;
+
+export function getServerMochaToBackground() {
+  if (!_serverMochaToBackground) {
+    startServerMochaToBackground();
+  }
+
+  return _serverMochaToBackground;
+}
+
+function startServerMochaToBackground() {
+  _serverMochaToBackground = new ServerMochaToBackground();
+}

test/e2e/background-socket/socket-background-to-mocha.ts

@@ -0,0 +1,136 @@
+import log from 'loglevel';
+import { WindowProperties } from './window-handles';
+
+/**
+ * This singleton class runs on the Extension background script (service worker in MV3).
+ * It's used to communicate from the Extension background script to the Mocha/Selenium test.
+ * The main advantage is that it can call chrome.tabs.query().
+ * We had hoped it would be able to call chrome.tabs.highlight(), but Selenium doesn't see the tab change.
+ */
+class SocketBackgroundToMocha {
+  private client: WebSocket;
+
+  constructor() {
+    this.client = new WebSocket('ws://localhost:8111');
+
+    this.client.onopen = () => log.debug('WebSocket connection opened');
+
+    this.client.onmessage = (event: MessageEvent) =>
+      this.receivedMessage(event.data);
+
+    this.client.onclose = () => log.debug('WebSocket connection closed');
+
+    this.client.onerror = (error) => log.error('WebSocket error:', error);
+  }
+
+  /**
+   * Waits until a window with the given property is open.
+   * delayStep = 200ms, timeout = 10s
+   *
+   * You can think of this kind of like a template function:
+   * If `property` is `title`, then this becomes `waitUntilWindowWithTitle`
+   * If `property` is `url`, then this becomes `waitUntilWindowWithUrl`
+   * Remember that `a[property]` becomes `a.title` or `a.url`
+   *
+   * @param property - 'title' or 'url'
+   * @param value - The value we're searching for and want to wait for
+   * @returns The handle of the window tab with the given property value
+   */
+  async waitUntilWindowWithProperty(property: WindowProperties, value: string) {
+    let tabs: chrome.tabs.Tab[] = [];
+    const delayStep = 200;
+    const timeout = 10000;
+
+    for (
+      let timeElapsed = 0;
+      timeElapsed <= timeout;
+      timeElapsed += delayStep
+    ) {
+      tabs = await this.queryTabs({});
+
+      const index = tabs.findIndex((a) => a[property] === value);
+
+      if (index !== -1) {
+        this.send({ command: 'openTabs', tabs: this.cleanTabs(tabs) });
+        return;
+      }
+
+      // wait for delayStep milliseconds
+      await new Promise((resolve) => setTimeout(resolve, delayStep));
+    }
+
+    // The window was not found at the end of the timeout
+    this.send({ command: 'notFound', tabs: this.cleanTabs(tabs) });
+  }
+
+  // This function exists to support both MV2 and MV3
+  private async queryTabs(queryInfo: object): Promise<chrome.tabs.Tab[]> {
+    if (
+      process.env.ENABLE_MV3 === 'true' ||
+      process.env.ENABLE_MV3 === undefined
+    ) {
+      // With MV3, chrome.tabs.query has an await form
+      return await chrome.tabs.query(queryInfo);
+    }
+
+    // With MV2, we have to wrap chrome.tabs.query in a Promise
+    return new Promise((resolve) => {
+      chrome.tabs.query(queryInfo, (tabs: chrome.tabs.Tab[]) => {
+        resolve(tabs);
+      });
+    });
+  }
+
+  // Clean up the tab data before sending them to the client
+  private cleanTabs(tabs: chrome.tabs.Tab[]): chrome.tabs.Tab[] {
+    return tabs.map((tab) => {
+      // This field can be very long, and is not needed
+      if (tab.favIconUrl && tab.favIconUrl.length > 40) {
+        tab.favIconUrl = undefined;
+      }
+
+      return tab;
+    });
+  }
+
+  // Send a message to the Mocha/Selenium test
+  send(message: string | object) {
+    if (typeof message === 'string') {
+      this.client.send(message);
+    } else {
+      this.client.send(JSON.stringify(message));
+    }
+  }
+
+  // Handle messages received from the Mocha/Selenium test
+  private async receivedMessage(message: string) {
+    const msg = JSON.parse(message);
+
+    log.debug('Received message:', msg);
+
+    if (msg.command === 'queryTabs') {
+      const tabs = await this.queryTabs({ title: msg.title });
+      log.debug('Sending tabs:', tabs);
+      this.send({ command: 'openTabs', tabs: this.cleanTabs(tabs) });
+    } else if (msg.command === 'waitUntilWindowWithProperty') {
+      this.waitUntilWindowWithProperty(msg.property, msg.value);
+    }
+  }
+}
+
+// Singleton setup below
+let _socketBackgroundToMocha: SocketBackgroundToMocha;
+
+export function getSocketBackgroundToMocha() {
+  if (!_socketBackgroundToMocha) {
+    startSocketBackgroundToMocha();
+  }
+
+  return _socketBackgroundToMocha;
+}
+
+function startSocketBackgroundToMocha() {
+  if (process.env.IN_TEST) {
+    _socketBackgroundToMocha = new SocketBackgroundToMocha();
+  }
+}