Fix more tests

Author: FrederikBolding

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

@@ -5233,8 +5233,8 @@ describe('SnapController', () => {
         },
       });
 
-      expect(rootMessenger.call).toHaveBeenCalledTimes(5);
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(options.messenger.call).toHaveBeenCalledTimes(5);
+      expect(options.messenger.call).toHaveBeenCalledWith(
         'ExecutionService:handleRpcRequest',
         MOCK_SNAP_ID,
         {
@@ -6437,7 +6437,7 @@ describe('SnapController', () => {
       });
       const [snapController] = getSnapControllerWithEES(snapControllerOptions);
 
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledWith(
         'PermissionController:grantPermissions',
         {
           approvedPermissions: {
@@ -6453,7 +6453,7 @@ describe('SnapController', () => {
         },
       );
 
-      expect(rootMessenger.publish).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.publish).toHaveBeenCalledWith(
         'SnapController:snapInstalled',
         getTruncatedSnap(),
         METAMASK_ORIGIN,
@@ -6537,12 +6537,12 @@ describe('SnapController', () => {
         () => MOCK_SNAP_PERMISSIONS,
       );
 
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledWith(
         'PermissionController:grantPermissions',
         { approvedPermissions, subject: { origin: 'npm:filsnap' } },
       );
 
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledWith(
         'PermissionController:grantPermissions',
         {
           approvedPermissions,
@@ -6586,7 +6586,7 @@ describe('SnapController', () => {
       });
       const [snapController] = getSnapControllerWithEES(snapControllerOptions);
 
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledWith(
         'PermissionController:grantPermissions',
         {
           approvedPermissions: {
@@ -6666,14 +6666,14 @@ describe('SnapController', () => {
       });
       const [snapController] = getSnapControllerWithEES(snapControllerOptions);
 
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledWith(
         'PermissionController:revokePermissions',
         {
           [MOCK_SNAP_ID]: ['snap_dialog'],
         },
       );
 
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledWith(
         'PermissionController:grantPermissions',
         {
           approvedPermissions: {
@@ -6684,7 +6684,7 @@ describe('SnapController', () => {
         },
       );
 
-      expect(rootMessenger.publish).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.publish).toHaveBeenCalledWith(
         'SnapController:snapUpdated',
         getTruncatedSnap({
           version: '1.2.3',
@@ -6753,7 +6753,7 @@ describe('SnapController', () => {
       });
       const [snapController] = getSnapControllerWithEES(snapControllerOptions);
 
-      expect(rootMessenger.call).toHaveBeenCalledTimes(0);
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledTimes(0);
 
       snapController.destroy();
     });
@@ -6803,7 +6803,7 @@ describe('SnapController', () => {
       });
       const [snapController] = getSnapControllerWithEES(snapControllerOptions);
 
-      expect(rootMessenger.call).toHaveBeenCalledWith(
+      expect(snapControllerOptions.messenger.call).toHaveBeenCalledWith(
         'PermissionController:grantPermissions',
         {
           approvedPermissions: {

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

@@ -20,14 +20,22 @@ export const getNodeEESMessenger = (
     ExecutionServiceActions,
     ExecutionServiceEvents
   >,
-) =>
-  new Messenger<
+) => {
+  const executionServiceMessenger = new Messenger<
     'ExecutionService',
     ExecutionServiceActions,
     ExecutionServiceEvents,
     any
   >({ namespace: 'ExecutionService', parent: messenger });
 
+  messenger.unregisterActionHandler('ExecutionService:handleRpcRequest');
+  messenger.unregisterActionHandler('ExecutionService:executeSnap');
+  messenger.unregisterActionHandler('ExecutionService:terminateSnap');
+  messenger.unregisterActionHandler('ExecutionService:terminateAllSnaps');
+
+  return executionServiceMessenger;
+};
+
 export const getNodeEES = (
   messenger: ReturnType<typeof getNodeEESMessenger>,
   setupSnapProvider?: SetupSnapProvider,

packages/snaps-simulation/src/simulation.ts

@@ -541,6 +541,8 @@ export function registerActions(
   options: SimulationOptions,
   snapId: SnapId,
 ) {
+  // TODO: Undo the changes in this file once you can unregister/register globally for tests
+
   controllerMessenger._internalRegisterDelegatedActionHandler(
     'PhishingController:testOrigin',
     () => ({ result: false, type: PhishingDetectorResultType.All }),

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

@@ -2,6 +2,7 @@ import type {
   ActionConstraint,
   ActionHandler,
   EventConstraint,
+  ExtractEventPayload,
 } from '@metamask/messenger';
 import { Messenger } from '@metamask/messenger';
 
@@ -24,7 +25,23 @@ 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);
+  }
 }