fix: Use latest DecryptMessageManager EncryptMessageManager to extend controllers

app/scripts/controllers/decrypt-message.test.ts

@@ -2,6 +2,7 @@ import {
   DecryptMessageManager,
   DecryptMessageParams,
 } from '@metamask/message-manager';
+import type { DecryptMessageManagerMessenger } from '@metamask/message-manager';
 import { MetaMetricsEventCategory } from '../../../shared/constants/metametrics';
 import DecryptMessageController, {
   DecryptMessageControllerMessenger,
@@ -36,11 +37,15 @@ const createMessengerMock = () =>
   ({
     registerActionHandler: jest.fn(),
     registerInitialEventPayload: jest.fn(),
+    subscribe: jest.fn(),
     publish: jest.fn(),
     call: jest.fn(),
-    // TODO: Replace `any` with type
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  } as any as jest.Mocked<DecryptMessageControllerMessenger>);
+  } as unknown as jest.Mocked<DecryptMessageControllerMessenger>);
+
+const createManagerMessengerMock = () =>
+  ({
+    subscribe: jest.fn(),
+  } as unknown as jest.Mocked<DecryptMessageManagerMessenger>);
 
 const createDecryptMessageManagerMock = <T>() =>
   ({
@@ -64,20 +69,14 @@ const createDecryptMessageManagerMock = <T>() =>
   } as any as jest.Mocked<T>);
 
 describe('DecryptMessageController', () => {
-  class MockDecryptMessageController extends DecryptMessageController {
-    // update is protected, so we expose it for typechecking here
-    public update(callback: Parameters<DecryptMessageController['update']>[0]) {
-      return super.update(callback);
-    }
-  }
-
-  let decryptMessageController: MockDecryptMessageController;
+  let decryptMessageController: DecryptMessageController;
 
   const decryptMessageManagerConstructorMock =
     DecryptMessageManager as jest.MockedClass<typeof DecryptMessageManager>;
   const getStateMock = jest.fn();
   const keyringControllerMock = createKeyringControllerMock();
   const messengerMock = createMessengerMock();
+  const managerMessengerMock = createManagerMessengerMock();
   const metricsEventMock = jest.fn();
 
   const decryptMessageManagerMock =
@@ -105,7 +104,7 @@ describe('DecryptMessageController', () => {
       decryptMessageManagerMock,
     );
 
-    decryptMessageController = new MockDecryptMessageController({
+    decryptMessageController = new DecryptMessageController({
       // TODO: Replace `any` with type
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       getState: getStateMock as any,
@@ -118,6 +117,7 @@ describe('DecryptMessageController', () => {
       // TODO: Replace `any` with type
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       metricsEvent: metricsEventMock as any,
+      managerMessenger: managerMessengerMock,
     } as DecryptMessageControllerOptions);
   });
 
@@ -127,23 +127,10 @@ describe('DecryptMessageController', () => {
   });
 
   it('should reset state', () => {
-    decryptMessageController.update(() => ({
-      unapprovedDecryptMsgs: {
-        [messageIdMock]: messageMock,
-        // TODO: Replace `any` with type
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      } as any,
-      unapprovedDecryptMsgCount: 1,
-    }));
     decryptMessageController.resetState();
     expect(decryptMessageController.state).toStrictEqual(getDefaultState());
   });
 
-  it('should clear unapproved messages', () => {
-    decryptMessageController.clearUnapproved();
-    expect(decryptMessageController.state).toStrictEqual(getDefaultState());
-    expect(decryptMessageManagerMock.update).toBeCalledTimes(1);
-  });
   it('should add unapproved messages', async () => {
     await decryptMessageController.newRequestDecryptMessage(
       messageMock,

app/scripts/controllers/decrypt-message.ts

@@ -2,7 +2,6 @@ import EventEmitter from 'events';
 import log from 'loglevel';
 import {
   AbstractMessage,
-  AbstractMessageManager,
   AbstractMessageParams,
   AbstractMessageParamsMetamask,
   MessageManagerState,
@@ -11,6 +10,10 @@ import {
   DecryptMessageParams,
   DecryptMessageParamsMetamask,
 } from '@metamask/message-manager';
+import type {
+  DecryptMessageManagerMessenger,
+  DecryptMessageManagerState,
+} from '@metamask/message-manager';
 import {
   BaseController,
   RestrictedControllerMessenger,
@@ -86,38 +89,47 @@ export type DecryptMessageControllerState = {
   unapprovedDecryptMsgCount: number;
 };
 
-export type GetDecryptMessageState = {
+export type GetDecryptMessageControllerState = {
   type: `${typeof controllerName}:getState`;
   handler: () => DecryptMessageControllerState;
 };
 
-export type DecryptMessageStateChange = {
+export type DecryptMessageControllerStateChange = {
   type: `${typeof controllerName}:stateChange`;
   payload: [DecryptMessageControllerState, Patch[]];
 };
 
-export type DecryptMessageControllerActions = GetDecryptMessageState;
+export type DecryptMessageControllerActions = GetDecryptMessageControllerState;
 
-export type DecryptMessageControllerEvents = DecryptMessageStateChange;
+export type DecryptMessageControllerEvents =
+  DecryptMessageControllerStateChange;
 
 type AllowedActions =
   | AddApprovalRequest
   | AcceptRequest
   | RejectRequest
   | KeyringControllerDecryptMessageAction;
 
+type DecryptMessageManagerStateChangeEvent = {
+  type: `DecryptMessageManager:stateChange`;
+  payload: [DecryptMessageManagerState, Patch[]];
+};
+
+type AllowedEvents = DecryptMessageManagerStateChangeEvent;
+
 export type DecryptMessageControllerMessenger = RestrictedControllerMessenger<
   typeof controllerName,
   DecryptMessageControllerActions | AllowedActions,
-  DecryptMessageControllerEvents,
+  DecryptMessageControllerEvents | AllowedEvents,
   AllowedActions['type'],
-  never
+  AllowedEvents['type']
 >;
 
 export type DecryptMessageControllerOptions = {
   // TODO: Replace `any` with type
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   getState: () => any;
+  managerMessenger: DecryptMessageManagerMessenger;
   messenger: DecryptMessageControllerMessenger;
   // TODO: Replace `any` with type
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -151,11 +163,13 @@ export default class DecryptMessageController extends BaseController<
    * @param options.getState - Callback to retrieve all user state.
    * @param options.messenger - A reference to the messaging system.
    * @param options.metricsEvent - A function for emitting a metric event.
+   * @param options.managerMessenger - A reference to the messenger need by the message manager.
    */
   constructor({
     getState,
     metricsEvent,
     messenger,
+    managerMessenger,
   }: DecryptMessageControllerOptions) {
     super({
       metadata: stateMetadata,
@@ -168,12 +182,11 @@ export default class DecryptMessageController extends BaseController<
 
     this.hub = new EventEmitter();
 
-    this._decryptMessageManager = new DecryptMessageManager(
-      undefined,
-      undefined,
-      undefined,
-      ['decrypted'],
-    );
+    this._decryptMessageManager = new DecryptMessageManager({
+      messenger: managerMessenger,
+      name: 'DecryptMessageManager',
+      additionalFinishStatuses: ['decrypted'],
+    });
 
     this._decryptMessageManager.hub.on('updateBadge', () => {
       this.hub.emit('updateBadge');
@@ -187,7 +200,7 @@ export default class DecryptMessageController extends BaseController<
     );
 
     this._subscribeToMessageState(
-      this._decryptMessageManager,
+      messenger,
       (state, newMessages, messageCount) => {
         state.unapprovedDecryptMsgs = newMessages;
         state.unapprovedDecryptMsgCount = messageCount;
@@ -215,10 +228,7 @@ export default class DecryptMessageController extends BaseController<
    * Clears all unapproved messages from memory.
    */
   clearUnapproved() {
-    this._decryptMessageManager.update({
-      unapprovedMessages: {},
-      unapprovedMessagesCount: 0,
-    });
+    this._decryptMessageManager.clearUnapprovedMessages();
   }
 
   /**
@@ -331,11 +341,7 @@ export default class DecryptMessageController extends BaseController<
   }
 
   private _cancelAbstractMessage(
-    messageManager: AbstractMessageManager<
-      AbstractMessage,
-      AbstractMessageParams,
-      AbstractMessageParamsMetamask
-    >,
+    messageManager: DecryptMessageManager,
     messageId: string,
     reason?: string,
   ) {
@@ -356,27 +362,26 @@ export default class DecryptMessageController extends BaseController<
   }
 
   private _subscribeToMessageState(
-    messageManager: AbstractMessageManager<
-      AbstractMessage,
-      AbstractMessageParams,
-      AbstractMessageParamsMetamask
-    >,
+    controllerMessenger: DecryptMessageControllerMessenger,
     updateState: (
       state: DecryptMessageControllerState,
       newMessages: Record<string, StateMessage>,
       messageCount: number,
     ) => void,
   ) {
-    messageManager.subscribe((state: MessageManagerState<AbstractMessage>) => {
-      const newMessages = this._migrateMessages(
-        // TODO: Replace `any` with type
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        state.unapprovedMessages as any,
-      );
-      this.update((draftState) => {
-        updateState(draftState, newMessages, state.unapprovedMessagesCount);
-      });
-    });
+    controllerMessenger.subscribe(
+      'DecryptMessageManager:stateChange',
+      (state: MessageManagerState<AbstractMessage>) => {
+        const newMessages = this._migrateMessages(
+          // TODO: Replace `any` with type
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          state.unapprovedMessages as any,
+        );
+        this.update((draftState) => {
+          updateState(draftState, newMessages, state.unapprovedMessagesCount);
+        });
+      },
+    );
   }
 
   private _migrateMessages(

app/scripts/controllers/encryption-public-key.test.ts

@@ -2,6 +2,7 @@ import {
   EncryptionPublicKeyManager,
   AbstractMessage,
   OriginalRequest,
+  EncryptionPublicKeyManagerMessenger,
 } from '@metamask/message-manager';
 import { KeyringType } from '../../../shared/constants/keyring';
 import { MetaMetricsEventCategory } from '../../../shared/constants/metametrics';
@@ -34,20 +35,7 @@ const messageMock = {
   status: 'unapproved',
   type: 'testType',
   rawSig: undefined,
-  // TODO: Replace `any` with type
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-} as any as AbstractMessage;
-
-const coreMessageMock = {
-  ...messageMock,
-  messageParams: messageParamsMock,
-};
-
-const stateMessageMock = {
-  ...messageMock,
-  msgParams: addressMock,
-  origin: messageParamsMock.origin,
-};
+} as unknown as AbstractMessage;
 
 const requestMock = {
   origin: 'http://test2.com',
@@ -57,11 +45,15 @@ const createMessengerMock = () =>
   ({
     registerActionHandler: jest.fn(),
     publish: jest.fn(),
+    subscribe: jest.fn(),
     call: jest.fn(),
     registerInitialEventPayload: jest.fn(),
-    // TODO: Replace `any` with type
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  } as any as jest.Mocked<EncryptionPublicKeyControllerMessenger>);
+  } as unknown as jest.Mocked<EncryptionPublicKeyControllerMessenger>);
+
+const createManagerMessengerMock = () =>
+  ({
+    subscribe: jest.fn(),
+  } as unknown as jest.Mocked<EncryptionPublicKeyManagerMessenger>);
 
 const createEncryptionPublicKeyManagerMock = <T>() =>
   ({
@@ -76,9 +68,7 @@ const createEncryptionPublicKeyManagerMock = <T>() =>
     hub: {
       on: jest.fn(),
     },
-    // TODO: Replace `any` with type
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  } as any as jest.Mocked<T>);
+  } as unknown as jest.Mocked<T>);
 
 describe('EncryptionPublicKeyController', () => {
   let encryptionPublicKeyController: EncryptionPublicKeyController;
@@ -90,6 +80,7 @@ describe('EncryptionPublicKeyController', () => {
   const encryptionPublicKeyManagerMock =
     createEncryptionPublicKeyManagerMock<EncryptionPublicKeyManager>();
   const messengerMock = createMessengerMock();
+  const managerMessengerMock = createManagerMessengerMock();
   const getEncryptionPublicKeyMock = jest.fn();
   const getAccountKeyringTypeMock = jest.fn();
   const getStateMock = jest.fn();
@@ -118,6 +109,7 @@ describe('EncryptionPublicKeyController', () => {
       // TODO: Replace `any` with type
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       metricsEvent: metricsEventMock as any,
+      managerMessenger: managerMessengerMock,
     } as EncryptionPublicKeyControllerOptions);
   });
 
@@ -198,22 +190,6 @@ describe('EncryptionPublicKeyController', () => {
     });
   });
 
-  describe('clearUnapproved', () => {
-    it('resets state in all message managers', () => {
-      encryptionPublicKeyController.clearUnapproved();
-
-      const defaultState = {
-        unapprovedMessages: {},
-        unapprovedMessagesCount: 0,
-      };
-
-      expect(encryptionPublicKeyManagerMock.update).toHaveBeenCalledTimes(1);
-      expect(encryptionPublicKeyManagerMock.update).toHaveBeenCalledWith(
-        defaultState,
-      );
-    });
-  });
-
   describe('newRequestEncryptionPublicKey', () => {
     // @ts-expect-error This function is missing from the Mocha type definitions
     it.each([
@@ -407,23 +383,5 @@ describe('EncryptionPublicKeyController', () => {
         true,
       );
     });
-
-    it('updates state on EncryptionPublicKeyManager state change', async () => {
-      await encryptionPublicKeyManagerMock.subscribe.mock.calls[0][0]({
-        // TODO: Replace `any` with type
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        unapprovedMessages: { [messageIdMock]: coreMessageMock as any },
-        unapprovedMessagesCount: 3,
-      });
-
-      expect(encryptionPublicKeyController.state).toEqual({
-        unapprovedEncryptionPublicKeyMsgs: {
-          // TODO: Replace `any` with type
-          // eslint-disable-next-line @typescript-eslint/no-explicit-any
-          [messageIdMock]: stateMessageMock as any,
-        },
-        unapprovedEncryptionPublicKeyMsgCount: 3,
-      });
-    });
   });
 });