Subscribe to event fired by KeyringController on account removal (#20478)

app/scripts/controllers/detect-tokens.test.js

@@ -226,6 +226,7 @@ describe('DetectTokensController', function () {
       onInfuraIsBlocked: sinon.stub(),
       onInfuraIsUnblocked: sinon.stub(),
       networkConfigurations: {},
+      onAccountRemoved: sinon.stub(),
     });
     preferences.setAddresses([
       '0x7e57e2',

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

@@ -52,6 +52,7 @@ describe('MMIController', function () {
         initState: {},
         onInfuraIsBlocked: jest.fn(),
         onInfuraIsUnblocked: jest.fn(),
+        onAccountRemoved: jest.fn(),
         provider: {},
         networkConfigurations: {},
       }),

app/scripts/controllers/preferences.js

@@ -117,6 +117,9 @@ export default class PreferencesController {
 
     this._subscribeToInfuraAvailability();
 
+    // subscribe to account removal
+    opts.onAccountRemoved((address) => this.removeAddress(address));
+
     global.setPreference = (key, value) => {
       return this.setFeatureFlag(key, value);
     };

app/scripts/controllers/preferences.test.js

@@ -42,6 +42,7 @@ describe('preferences controller', () => {
       tokenListController,
       onInfuraIsBlocked: jest.fn(),
       onInfuraIsUnblocked: jest.fn(),
+      onAccountRemoved: jest.fn(),
       networkConfigurations: NETWORK_CONFIGURATION_DATA,
     });
   });
@@ -109,6 +110,65 @@ describe('preferences controller', () => {
     });
   });
 
+  describe('onAccountRemoved', () => {
+    it('should remove an address from state', () => {
+      const testAddress = '0xda22le';
+      let accountRemovedListener;
+      const onAccountRemoved = (callback) => {
+        accountRemovedListener = callback;
+      };
+      preferencesController = new PreferencesController({
+        initLangCode: 'en_US',
+        tokenListController,
+        onInfuraIsBlocked: jest.fn(),
+        onInfuraIsUnblocked: jest.fn(),
+        initState: {
+          identities: {
+            [testAddress]: {
+              name: 'Account 1',
+              address: testAddress,
+            },
+          },
+        },
+        onAccountRemoved,
+        networkConfigurations: NETWORK_CONFIGURATION_DATA,
+      });
+
+      accountRemovedListener(testAddress);
+
+      expect(
+        preferencesController.store.getState().identities['0xda22le'],
+      ).toStrictEqual(undefined);
+    });
+
+    it('should throw an error if address not found', () => {
+      const testAddress = '0xda22le';
+      let accountRemovedListener;
+      const onAccountRemoved = (callback) => {
+        accountRemovedListener = callback;
+      };
+      preferencesController = new PreferencesController({
+        initLangCode: 'en_US',
+        tokenListController,
+        onInfuraIsBlocked: jest.fn(),
+        onInfuraIsUnblocked: jest.fn(),
+        initState: {
+          identities: {
+            '0x7e57e2': {
+              name: 'Account 1',
+              address: '0x7e57e2',
+            },
+          },
+        },
+        onAccountRemoved,
+        networkConfigurations: NETWORK_CONFIGURATION_DATA,
+      });
+      expect(() => {
+        accountRemovedListener(testAddress);
+      }).toThrow(`${testAddress} can't be deleted cause it was not found`);
+    });
+  });
+
   describe('removeAddress', () => {
     it('should remove an address from state', () => {
       preferencesController.setAddresses(['0xda22le', '0x7e57e2']);

app/scripts/lib/account-tracker.js

@@ -56,6 +56,7 @@ export default class AccountTracker {
    * @param {object} opts.blockTracker - A block tracker, which emits events for each new block
    * @param {Function} opts.getCurrentChainId - A function that returns the `chainId` for the current global network
    * @param {Function} opts.getNetworkIdentifier - A function that returns the current network
+   * @param {Function} opts.onAccountRemoved - Allows subscribing to keyring controller accountRemoved event
    */
   constructor(opts = {}) {
     const initState = {
@@ -83,6 +84,9 @@ export default class AccountTracker {
     this.preferencesController = opts.preferencesController;
     this.onboardingController = opts.onboardingController;
 
+    // subscribe to account removal
+    opts.onAccountRemoved((address) => this.removeAccount([address]));
+
     this.onboardingController.store.subscribe(
       previousValueComparator(async (prevState, currState) => {
         const { completedOnboarding: prevCompletedOnboarding } = prevState;

app/scripts/lib/account-tracker.test.js

@@ -70,6 +70,7 @@ describe('Account Tracker', () => {
           getState: noop,
         },
       },
+      onAccountRemoved: jest.fn(),
     });
   });
 
@@ -140,6 +141,38 @@ describe('Account Tracker', () => {
     });
   });
 
+  describe('onAccountRemoved', () => {
+    it('should remove an account from state', () => {
+      let accountRemovedListener;
+      const onAccountRemoved = (callback) => {
+        accountRemovedListener = callback;
+      };
+      accountTracker = new AccountTracker({
+        provider,
+        blockTracker: blockTrackerStub,
+        preferencesController: {
+          store: {
+            getState: () => ({
+              useMultiAccountBalanceChecker,
+            }),
+            subscribe: noop,
+          },
+        },
+        onboardingController: {
+          store: {
+            subscribe: noop,
+            getState: noop,
+          },
+        },
+        onAccountRemoved,
+      });
+      accountRemovedListener(VALID_ADDRESS);
+      expect(
+        accountTracker.store.getState().accounts[VALID_ADDRESS],
+      ).toStrictEqual(undefined);
+    });
+  });
+
   describe('_updateAccountsViaBalanceChecker', () => {
     it('should update the passed address account balance, and set other balances to null, if useMultiAccountBalanceChecker is false', async () => {
       useMultiAccountBalanceChecker = true;

app/scripts/metamask-controller.js

@@ -412,6 +412,10 @@ export default class MetamaskController extends EventEmitter {
         networkControllerMessenger,
         'NetworkController:infuraIsUnblocked',
       ),
+      onAccountRemoved: this.controllerMessenger.subscribe.bind(
+        this.controllerMessenger,
+        'KeyringController:accountRemoved',
+      ),
       tokenListController: this.tokenListController,
       provider: this.provider,
       networkConfigurations: this.networkController.state.networkConfigurations,
@@ -757,6 +761,10 @@ export default class MetamaskController extends EventEmitter {
         initState.AccountTracker?.accounts
           ? { accounts: initState.AccountTracker.accounts }
           : { accounts: {} },
+      onAccountRemoved: this.controllerMessenger.subscribe.bind(
+        this.controllerMessenger,
+        'KeyringController:accountRemoved',
+      ),
     });
 
     // start and stop polling for balances based on activeControllerConnections
@@ -3551,10 +3559,6 @@ export default class MetamaskController extends EventEmitter {
   async removeAccount(address) {
     // Remove all associated permissions
     this.removeAllAccountPermissions(address);
-    // Remove account from the preferences controller
-    this.preferencesController.removeAddress(address);
-    // Remove account from the account tracker controller
-    this.accountTracker.removeAccount([address]);
 
     ///: BEGIN:ONLY_INCLUDE_IN(build-mmi)
     this.custodyController.removeAccount(address);

app/scripts/metamask-controller.test.js

@@ -944,8 +944,6 @@ describe('MetaMaskController', function () {
           getAccounts: sinon.stub().returns(Promise.resolve([])),
           destroy: sinon.stub(),
         };
-        sinon.stub(metamaskController.preferencesController, 'removeAddress');
-        sinon.stub(metamaskController.accountTracker, 'removeAccount');
         sinon.stub(metamaskController.keyringController, 'removeAccount');
         sinon.stub(metamaskController, 'removeAllAccountPermissions');
         sinon
@@ -960,28 +958,12 @@ describe('MetaMaskController', function () {
 
       afterEach(function () {
         metamaskController.keyringController.removeAccount.restore();
-        metamaskController.accountTracker.removeAccount.restore();
-        metamaskController.preferencesController.removeAddress.restore();
         metamaskController.removeAllAccountPermissions.restore();
 
         mockKeyring.getAccounts.resetHistory();
         mockKeyring.destroy.resetHistory();
       });
 
-      it('should call preferencesController.removeAddress', async function () {
-        assert(
-          metamaskController.preferencesController.removeAddress.calledWith(
-            addressToRemove,
-          ),
-        );
-      });
-      it('should call accountTracker.removeAccount', async function () {
-        assert(
-          metamaskController.accountTracker.removeAccount.calledWith([
-            addressToRemove,
-          ]),
-        );
-      });
       it('should call keyringController.removeAccount', async function () {
         assert(
           metamaskController.keyringController.removeAccount.calledWith(