Fix type errors, revert some changes, and get more tests passing

Author: Mrtenz

packages/snaps-controllers/src/snaps/SnapController.test.tsx

@@ -10295,12 +10295,13 @@ describe('SnapController', () => {
       expect(updatedSnap.version).toStrictEqual(updateVersion);
       expect(updatedSnap.preinstalled).toBe(true);
 
-      expect(rootMessenger.call).toHaveBeenNthCalledWith(
+      expect(messenger.call).toHaveBeenNthCalledWith(
         7,
         'PermissionController:revokePermissions',
         { [snapId]: [SnapEndowments.Rpc, SnapEndowments.LifecycleHooks] },
       );
-      expect(rootMessenger.call).toHaveBeenNthCalledWith(
+
+      expect(messenger.call).toHaveBeenNthCalledWith(
         8,
         'PermissionController:grantPermissions',
         {

…controllers/src/test-utils/controller.ts …ontrollers/src/test-utils/controller.tsxpackages/snaps-controllers/src/test-utils/controller.ts renamed to packages/snaps-controllers/src/test-utils/controller.tsx packages/snaps-controllers/src/test-utils/controller.ts renamed to packages/snaps-controllers/src/test-utils/controller.tsx

@@ -1,38 +1,39 @@
 import type { ApprovalRequest } from '@metamask/approval-controller';
 import {
-  encryptWithKey,
   decryptWithKey,
-  keyFromPassword,
-  importKey,
+  encryptWithKey,
   exportKey,
   generateSalt,
+  importKey,
   isVaultUpdated,
+  keyFromPassword,
 } from '@metamask/browser-passworder';
 import { Messenger } from '@metamask/messenger';
 import type {
+  Caveat,
   PermissionConstraint,
+  SubjectMetadata,
   SubjectPermissions,
   ValidPermission,
-  Caveat,
-  SubjectMetadata,
 } from '@metamask/permission-controller';
 import { SubjectType } from '@metamask/permission-controller';
+import { PhishingDetectorResultType } from '@metamask/phishing-controller';
 import { providerErrors } from '@metamask/rpc-errors';
 import {
   SnapEndowments,
   WALLET_SNAP_PERMISSION_KEY,
 } from '@metamask/snaps-rpc-methods';
 import type { SnapId } from '@metamask/snaps-sdk';
-import { text } from '@metamask/snaps-sdk';
+import { Text } from '@metamask/snaps-sdk/jsx';
 import { SnapCaveatType } from '@metamask/snaps-utils';
 import {
-  MockControllerMessenger,
   getPersistedSnapObject,
+  getSnapObject,
   MOCK_LOCAL_SNAP_ID,
   MOCK_ORIGIN,
   MOCK_SNAP_ID,
+  MockControllerMessenger,
   TEST_SECRET_RECOVERY_PHRASE_SEED_BYTES,
-  getSnapObject,
 } from '@metamask/snaps-utils/test-utils';
 import type { Json } from '@metamask/utils';
 
@@ -58,7 +59,6 @@ import type {
   MultichainRouterAllowedActions,
   MultichainRouterEvents,
 } from '../multichain';
-import { SnapController } from '../snaps';
 import type {
   AllowedActions,
   AllowedEvents,
@@ -69,6 +69,7 @@ import type {
   SnapsRegistryActions,
   SnapsRegistryEvents,
 } from '../snaps';
+import { SnapController } from '../snaps';
 import type { KeyDerivationOptions } from '../types';
 import type {
   WebSocketServiceActions,
@@ -441,9 +442,10 @@ export const getControllerMessenger = (registry = new MockSnapsRegistry()) => {
       if (id !== MOCK_INTERFACE_ID) {
         throw new Error(`Interface with id '${id}' not found.`);
       }
+
       return {
         snapId,
-        content: text('foo bar'),
+        content: <Text>console.log('hello world');</Text>,
         state: {},
         context: null,
       } as StoredInterface;
@@ -621,6 +623,7 @@ export const getSnapControllerWithEES = (
   service?: ReturnType<typeof getNodeEES>,
 ) => {
   const _service =
+    // @ts-expect-error: TODO: Investigate type mismatch.
     service ?? getNodeEES(getNodeEESMessenger(options.rootMessenger));
 
   const controller = new SnapController(options);
@@ -773,7 +776,7 @@ export const getRestrictedSnapInterfaceControllerMessenger = (
   if (mocked) {
     messenger.registerActionHandler('PhishingController:testOrigin', () => ({
       result: false,
-      type: 'all',
+      type: PhishingDetectorResultType.All,
     }));
 
     messenger.registerActionHandler(
@@ -949,7 +952,7 @@ export const getRestrictedWebSocketServiceMessenger = (
 ) => {
   const controllerMessenger = new Messenger<
     'WebSocketService',
-    WebSocketServiceAllowedActions,
+    WebSocketServiceActions | WebSocketServiceAllowedActions,
     WebSocketServiceEvents,
     any
   >({ namespace: 'WebSocketService', parent: messenger });
@@ -964,5 +967,7 @@ export const getRestrictedWebSocketServiceMessenger = (
     messenger: controllerMessenger,
   });
 
+  jest.spyOn(controllerMessenger, 'call');
+
   return controllerMessenger;
 };

packages/snaps-controllers/src/test-utils/execution-environment.ts

@@ -65,6 +65,10 @@ export const getNodeEES = (
   });
 
 export class ExecutionEnvironmentStub implements ExecutionService {
+  name: 'ExecutionService' = 'ExecutionService' as const;
+
+  state = null;
+
   constructor(messenger: ReturnType<typeof getNodeEESMessenger>) {
     messenger.registerActionHandler(
       `ExecutionService:handleRpcRequest`,

packages/snaps-controllers/src/websocket/WebSocketService.test.ts

@@ -18,9 +18,9 @@ const MOCK_WEBSOCKET_BROKEN_URI = 'wss://broken.metamask.io';
 class MockWebSocket {
   readonly #origin: string;
 
-  #closeListener: EventListener;
+  #closeListener?: EventListener;
 
-  #messageListener: EventListener;
+  #messageListener?: EventListener;
 
   constructor(url: string) {
     this.#origin = new URL(url).origin;
@@ -47,7 +47,7 @@ class MockWebSocket {
 
   send(data: string | Uint8Array) {
     if (data === 'Ping') {
-      this.#messageListener(
+      this.#messageListener?.(
         new MessageEvent('message', { origin: this.#origin, data: 'Pong' }),
       );
       return;
@@ -58,7 +58,7 @@ class MockWebSocket {
       isEqual(Array.from(data), Array.from(stringToBytes('Ping')))
     ) {
       const bytes = stringToBytes('Pong');
-      this.#messageListener(
+      this.#messageListener?.(
         new MessageEvent('message', {
           origin: this.#origin,
           data: bytes.buffer,
@@ -73,7 +73,7 @@ class MockWebSocket {
     event.code = code;
     event.reason = reason;
     event.wasClean = wasClean;
-    this.#closeListener(event);
+    this.#closeListener?.(event);
   }
 
   close() {
@@ -117,7 +117,7 @@ describe('WebSocketService', () => {
       'Ping',
     );
 
-    expect(rootMessenger.call).toHaveBeenNthCalledWith(
+    expect(messenger.call).toHaveBeenNthCalledWith(
       2,
       'SnapController:handleRequest',
       {
@@ -137,7 +137,7 @@ describe('WebSocketService', () => {
       },
     );
 
-    expect(rootMessenger.call).toHaveBeenNthCalledWith(
+    expect(messenger.call).toHaveBeenNthCalledWith(
       4,
       'SnapController:handleRequest',
       {
@@ -192,7 +192,7 @@ describe('WebSocketService', () => {
       Array.from(stringToBytes('Ping')),
     );
 
-    expect(rootMessenger.call).toHaveBeenNthCalledWith(
+    expect(messenger.call).toHaveBeenNthCalledWith(
       2,
       'SnapController:handleRequest',
       {
@@ -212,7 +212,7 @@ describe('WebSocketService', () => {
       },
     );
 
-    expect(rootMessenger.call).toHaveBeenNthCalledWith(
+    expect(messenger.call).toHaveBeenNthCalledWith(
       4,
       'SnapController:handleRequest',
       {
@@ -327,7 +327,7 @@ describe('WebSocketService', () => {
       messenger.call('WebSocketService:getAll', MOCK_SNAP_ID),
     ).toHaveLength(0);
 
-    expect(rootMessenger.call).toHaveBeenNthCalledWith(
+    expect(messenger.call).toHaveBeenNthCalledWith(
       5,
       'SnapController:handleRequest',
       {

packages/snaps-utils/src/test-utils/controller.ts

@@ -2,16 +2,16 @@ import type {
   ActionConstraint,
   ActionHandler,
   EventConstraint,
-  ExtractEventPayload,
+  MockAnyNamespace,
 } from '@metamask/messenger';
-import { Messenger } from '@metamask/messenger';
+import { MOCK_ANY_NAMESPACE, Messenger } from '@metamask/messenger';
 
 export class MockControllerMessenger<
   Action extends ActionConstraint,
   Event extends EventConstraint,
-> extends Messenger<'MockMessenger', Action, Event> {
+> extends Messenger<MockAnyNamespace, Action, Event> {
   constructor() {
-    super({ namespace: 'MockMessenger' });
+    super({ namespace: MOCK_ANY_NAMESPACE });
   }
 
   /**
@@ -25,23 +25,7 @@ export class MockControllerMessenger<
     actionType: ActionType,
     handler: ActionHandler<Action, ActionType>,
   ) {
-    // TODO: Undo this once you can unregister/register globally for tests
-    super._internalUnregisterDelegatedActionHandler(actionType);
-    super._internalRegisterDelegatedActionHandler(actionType, handler);
-  }
-
-  unregisterActionHandler<ActionType extends Action['type']>(
-    actionType: ActionType,
-  ) {
-    // TODO: Undo this once you can unregister/register globally for tests
-    super._internalUnregisterDelegatedActionHandler(actionType);
-  }
-
-  publish<EventType extends Event['type']>(
-    eventType: EventType,
-    ...payload: ExtractEventPayload<Event, EventType>
-  ): void {
-    // TODO: Undo this once you can publish globally for tests
-    super._internalPublishDelegated(eventType, ...payload);
+    super.unregisterActionHandler(actionType);
+    super.registerActionHandler(actionType, handler);
   }
 }