refactor: migrate ShieldController to @metamask/messenger

packages/shield-controller/CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- **BREAKING:** Use new `Messenger` from `@metamask/messenger` ([#6497](https://github.com/MetaMask/core/pull/6497))
+  - Previously, `ShieldController` accepted a `RestrictedMessenger` instance from `@metamask/base-controller`.
 - Bump `@metamask/base-controller` from `^8.2.0` to `^8.3.0` ([#6465](https://github.com/MetaMask/core/pull/6465))
 
 ## [0.1.2]

packages/shield-controller/package.json

@@ -48,6 +48,7 @@
   },
   "dependencies": {
     "@metamask/base-controller": "^8.3.0",
+    "@metamask/messenger": "^0.2.0",
     "@metamask/utils": "^11.4.2"
   },
   "devDependencies": {

packages/shield-controller/src/ShieldController.test.ts

@@ -21,7 +21,7 @@ function setup({
   transactionHistoryLimit?: number;
 } = {}) {
   const backend = createMockBackend();
-  const { messenger, baseMessenger } = createMockMessenger();
+  const { messenger, rootMessenger } = createMockMessenger();
 
   const controller = new ShieldController({
     backend,
@@ -33,23 +33,23 @@ function setup({
   return {
     controller,
     messenger,
-    baseMessenger,
+    rootMessenger,
     backend,
   };
 }
 
 describe('ShieldController', () => {
   describe('checkCoverage', () => {
     it('should trigger checkCoverage when a new transaction is added', async () => {
-      const { baseMessenger, backend } = setup();
+      const { rootMessenger, backend } = setup();
       const txMeta = generateMockTxMeta();
       const coverageResultReceived = new Promise<void>((resolve) => {
-        baseMessenger.subscribe(
+        rootMessenger.subscribe(
           'ShieldController:coverageResultReceived',
           (_coverageResult) => resolve(),
         );
       });
-      baseMessenger.publish(
+      rootMessenger.publish(
         'TransactionController:stateChange',
         { transactions: [txMeta] } as TransactionControllerState,
         undefined as never,
@@ -59,11 +59,11 @@ describe('ShieldController', () => {
     });
 
     it('should no longer trigger checkCoverage when controller is stopped', async () => {
-      const { controller, baseMessenger, backend } = setup();
+      const { controller, rootMessenger, backend } = setup();
       controller.stop();
       const txMeta = generateMockTxMeta();
       const coverageResultReceived = new Promise<void>((resolve, reject) => {
-        baseMessenger.subscribe(
+        rootMessenger.subscribe(
           'ShieldController:coverageResultReceived',
           (_coverageResult) => resolve(),
         );
@@ -72,7 +72,7 @@ describe('ShieldController', () => {
           100,
         );
       });
-      baseMessenger.publish(
+      rootMessenger.publish(
         'TransactionController:stateChange',
         { transactions: [txMeta] } as TransactionControllerState,
         undefined as never,
@@ -111,17 +111,17 @@ describe('ShieldController', () => {
     });
 
     it('should check coverage when a transaction is simulated', async () => {
-      const { baseMessenger, backend } = setup();
+      const { rootMessenger, backend } = setup();
       const txMeta = generateMockTxMeta();
       const coverageResultReceived = new Promise<void>((resolve) => {
-        baseMessenger.subscribe(
+        rootMessenger.subscribe(
           'ShieldController:coverageResultReceived',
           (_coverageResult) => resolve(),
         );
       });
 
       // Add transaction.
-      baseMessenger.publish(
+      rootMessenger.publish(
         'TransactionController:stateChange',
         { transactions: [txMeta] } as TransactionControllerState,
         undefined as never,
@@ -133,7 +133,7 @@ describe('ShieldController', () => {
       txMeta.simulationData = {
         tokenBalanceChanges: [],
       };
-      baseMessenger.publish(
+      rootMessenger.publish(
         'TransactionController:stateChange',
         { transactions: [txMeta] } as TransactionControllerState,
         undefined as never,

packages/shield-controller/src/ShieldController.ts

@@ -1,8 +1,9 @@
-import { BaseController } from '@metamask/base-controller';
-import type {
-  ControllerStateChangeEvent,
-  RestrictedMessenger,
-} from '@metamask/base-controller';
+import {
+  BaseController,
+  type ControllerGetStateAction,
+  type ControllerStateChangeEvent,
+} from '@metamask/base-controller/next';
+import type { Messenger } from '@metamask/messenger';
 import type {
   TransactionControllerStateChangeEvent,
   TransactionMeta,
@@ -47,6 +48,11 @@ export function getDefaultShieldControllerState(): ShieldControllerState {
   };
 }
 
+export type ShieldControllerGetStateAction = ControllerGetStateAction<
+  typeof controllerName,
+  ShieldControllerState
+>;
+
 export type ShieldControllerCheckCoverageAction = {
   type: `${typeof controllerName}:checkCoverage`;
   handler: ShieldController['checkCoverage'];
@@ -55,7 +61,9 @@ export type ShieldControllerCheckCoverageAction = {
 /**
  * The internal actions available to the ShieldController.
  */
-export type ShieldControllerActions = ShieldControllerCheckCoverageAction;
+export type ShieldControllerActions =
+  | ShieldControllerGetStateAction
+  | ShieldControllerCheckCoverageAction;
 
 export type ShieldControllerCoverageResultReceivedEvent = {
   type: `${typeof controllerName}:coverageResultReceived`;
@@ -74,11 +82,6 @@ export type ShieldControllerEvents =
   | ShieldControllerCoverageResultReceivedEvent
   | ShieldControllerStateChangeEvent;
 
-/**
- * The external actions available to the ShieldController.
- */
-type AllowedActions = never;
-
 /**
  * The external events available to the ShieldController.
  */
@@ -87,12 +90,10 @@ type AllowedEvents = TransactionControllerStateChangeEvent;
 /**
  * The messenger of the {@link ShieldController}.
  */
-export type ShieldControllerMessenger = RestrictedMessenger<
+export type ShieldControllerMessenger = Messenger<
   typeof controllerName,
-  ShieldControllerActions | AllowedActions,
-  ShieldControllerEvents | AllowedEvents,
-  AllowedActions['type'],
-  AllowedEvents['type']
+  ShieldControllerActions,
+  ShieldControllerEvents | AllowedEvents
 >;
 
 /**
@@ -160,15 +161,15 @@ export class ShieldController extends BaseController<
   }
 
   start() {
-    this.messagingSystem.subscribe(
+    this.messenger.subscribe(
       'TransactionController:stateChange',
       this.#transactionControllerStateChangeHandler,
       (state) => state.transactions,
     );
   }
 
   stop() {
-    this.messagingSystem.unsubscribe(
+    this.messenger.unsubscribe(
       'TransactionController:stateChange',
       this.#transactionControllerStateChangeHandler,
     );
@@ -211,7 +212,7 @@ export class ShieldController extends BaseController<
     const coverageResult = await this.#fetchCoverageResult(txMeta);
 
     // Publish coverage result
-    this.messagingSystem.publish(
+    this.messenger.publish(
       `${controllerName}:coverageResultReceived`,
       coverageResult,
     );

packages/shield-controller/tests/mocks/messenger.ts

@@ -1,34 +1,61 @@
-import { Messenger } from '@metamask/base-controller';
-
-import type {
-  ExtractAvailableAction,
-  ExtractAvailableEvent,
-} from '../../../base-controller/tests/helpers';
-import type { ShieldControllerActions } from '../../src';
 import {
-  type ShieldControllerEvents,
-  type ShieldControllerMessenger,
-} from '../../src';
+  Messenger,
+  MOCK_ANY_NAMESPACE,
+  type MessengerActions,
+  type MessengerEvents,
+  type MockAnyNamespace,
+} from '@metamask/messenger';
+
+import { type ShieldControllerMessenger } from '../../src';
 import { controllerName } from '../../src/constants';
 
+type AllShieldControllerActions = MessengerActions<ShieldControllerMessenger>;
+
+type AllShieldControllerEvents = MessengerEvents<ShieldControllerMessenger>;
+
+type RootMessenger = Messenger<
+  MockAnyNamespace,
+  AllShieldControllerActions,
+  AllShieldControllerEvents
+>;
+
+/**
+ * Constructs the root messenger.
+ *
+ * @returns A root messenger.
+ */
+function getRootMessenger(): RootMessenger {
+  return new Messenger({
+    namespace: MOCK_ANY_NAMESPACE,
+  });
+}
+
 /**
  * Create a mock messenger.
  *
  * @returns A mock messenger.
  */
-export function createMockMessenger() {
-  const baseMessenger = new Messenger<
-    ShieldControllerActions | ExtractAvailableAction<ShieldControllerMessenger>,
-    ShieldControllerEvents | ExtractAvailableEvent<ShieldControllerMessenger>
-  >();
-  const messenger = baseMessenger.getRestricted({
-    name: controllerName,
-    allowedActions: [],
-    allowedEvents: ['TransactionController:stateChange'],
+export function createMockMessenger(): {
+  rootMessenger: RootMessenger;
+  messenger: ShieldControllerMessenger;
+} {
+  const rootMessenger = getRootMessenger();
+  const messenger = new Messenger<
+    typeof controllerName,
+    AllShieldControllerActions,
+    AllShieldControllerEvents,
+    RootMessenger
+  >({
+    namespace: controllerName,
+    parent: rootMessenger,
+  });
+  rootMessenger.delegate({
+    messenger,
+    events: ['TransactionController:stateChange'],
   });
 
   return {
-    baseMessenger,
+    rootMessenger,
     messenger,
   };
 }

yarn.lock

@@ -4469,6 +4469,7 @@ __metadata:
     "@lavamoat/preinstall-always-fail": "npm:^2.1.0"
     "@metamask/auto-changelog": "npm:^3.4.4"
     "@metamask/base-controller": "npm:^8.3.0"
+    "@metamask/messenger": "npm:^0.2.0"
     "@metamask/transaction-controller": "npm:^60.2.0"
     "@metamask/utils": "npm:^11.4.2"
     "@ts-bridge/cli": "npm:^0.6.1"