changes based on the review.

Author: kanthesha

app/scripts/controllers/mmi-controller.test.ts

@@ -40,6 +40,21 @@ jest.mock('./permissions', () => ({
   }),
 }));
 
+export const createMockNetworkConfiguration = (
+  override?: Partial<NetworkConfiguration>,
+): NetworkConfiguration => {
+  return {
+    chainId: CHAIN_IDS.SEPOLIA,
+    blockExplorerUrls: [],
+    defaultRpcEndpointIndex: 0,
+    name: 'Mock Network',
+    nativeCurrency: 'MOCK TOKEN',
+    rpcEndpoints: [],
+    defaultBlockExplorerUrlIndex: 0,
+    ...override,
+  };
+};
+
 const mockAccount = {
   address: '0x758b8178a9A4B7206d1f648c4a77C515Cbac7001',
   id: 'mock-id',
@@ -210,6 +225,20 @@ describe('MMIController', function () {
       InfuraNetworkType['sepolia'],
     );
 
+    controllerMessenger.registerActionHandler(
+      'NetworkController:getNetworkClientById',
+      jest.fn().mockReturnValue({
+        configuration: {
+          chainId: CHAIN_IDS.SEPOLIA,
+        }
+      }),
+    );
+
+    controllerMessenger.registerActionHandler(
+      'NetworkController:getNetworkConfigurationByChainId',
+      jest.fn().mockReturnValue(createMockNetworkConfiguration()),
+    );
+
     const mmiControllerMessenger = controllerMessenger.getRestricted({
       name: 'MMIController',
       allowedActions: [
@@ -219,7 +248,9 @@ describe('MMIController', function () {
         'AccountsController:getSelectedAccount',
         'AccountsController:setSelectedAccount',
         'NetworkController:getState',
-        'NetworkController:setActiveNetwork'
+        'NetworkController:setActiveNetwork',
+        'NetworkController:getNetworkClientById',
+        'NetworkController:getNetworkConfigurationByChainId'
       ],
     });
 
@@ -508,9 +539,7 @@ describe('MMIController', function () {
       CUSTODIAN_TYPES['CUSTODIAN-TYPE'] = {
         keyringClass: { type: 'mock-keyring-class' },
       };
-      mmiController.messagingSystem.call = jest
-        .fn()
-        .mockReturnValue({ address: '0x1' });
+      jest.spyOn(ControllerMessenger.prototype, 'call').mockReturnValue({ address: '0x1' });
       mmiController.custodyController.getCustodyTypeByAddress = jest
         .fn()
         .mockReturnValue('custodian-type');
@@ -837,7 +866,7 @@ describe('MMIController', function () {
         '0x1',
       );
 
-      expect(selectedAccountSpy).toHaveBeenCalledTimes(1);
+      expect(selectedAccountSpy).toHaveBeenCalledTimes(4);
       const selectedAccount = accountsController.getSelectedAccount();
       expect(selectedAccount.id).toBe(mockAccount.id);
     });

app/scripts/controllers/mmi-controller.ts

@@ -42,9 +42,6 @@ import {
   ConnectionRequest,
   MMIControllerMessenger,
 } from '../../../shared/constants/mmi-controller';
-// TODO: Remove restricted import
-// eslint-disable-next-line import/no-restricted-paths
-import { getCurrentChainId } from '../../../ui/selectors';
 import MetaMetricsController from './metametrics';
 import { getPermissionBackgroundApiMethods } from './permissions';
 import AccountTrackerController from './account-tracker-controller';
@@ -866,15 +863,21 @@ export class MMIController {
       }
     }
 
-    const selectedChainId = getCurrentChainId({
-      metamask: this.#networkControllerState,
-    });
+    const { selectedNetworkClientId } = this.messagingSystem.call(
+      'NetworkController:getState',
+    );
+    const {
+      configuration: { chainId: selectedChainId },
+    } = this.messagingSystem.call(
+      'NetworkController:getNetworkClientById',
+      selectedNetworkClientId,
+    );
 
     if (selectedChainId !== toHex(chainId)) {
-      const networkConfiguration =
-        this.#networkControllerState.networkConfigurationsByChainId[
-          toHex(chainId)
-        ];
+      const networkConfiguration = this.messagingSystem.call(
+        'NetworkController:getNetworkConfigurationByChainId',
+        toHex(chainId),
+      );
 
       const { networkClientId } =
         networkConfiguration?.rpcEndpoints?.[

shared/constants/mmi-controller.ts

@@ -5,6 +5,7 @@ import { CustodyController } from '@metamask-institutional/custody-controller';
 import { SignatureController } from '@metamask/signature-controller';
 import {
   NetworkController,
+  NetworkControllerGetNetworkClientByIdAction,
   NetworkControllerGetStateAction,
   NetworkControllerSetActiveNetworkAction,
 } from '@metamask/network-controller';
@@ -31,6 +32,11 @@ import MetaMetricsController from '../../app/scripts/controllers/metametrics';
 // Unique name for the controller
 const controllerName = 'MMIController';
 
+type NetworkControllerGetNetworkConfigurationByChainId = {
+  type: `NetworkController:getNetworkConfigurationByChainId`;
+  handler: NetworkController['getNetworkConfigurationByChainId'];
+};
+
 /**
  * Actions that this controller is allowed to call.
  */
@@ -41,7 +47,9 @@ export type AllowedActions =
   | AccountsControllerGetSelectedAccountAction
   | AccountsControllerSetSelectedAccountAction
   | NetworkControllerGetStateAction
-  | NetworkControllerSetActiveNetworkAction;
+  | NetworkControllerSetActiveNetworkAction
+  | NetworkControllerGetNetworkClientByIdAction
+  | NetworkControllerGetNetworkConfigurationByChainId;
 
 /**
  * Messenger type for the {@link MMIController}.