fix: update handleAccountRemoved logic (#4322)

montelaidev · ccharly · web-flow · commit 5761ee77a9a3 · 2024-06-25T21:47:28.000+08:00
## Explanation This PR modifies the handleAccountRemoved function to update the selected account to the most recent one, if the account being removed is currently the selected account, all within the same update. ## References Related to MetaMask/metamask-mobile#9749 ## Changelog ### `@metamask/accounts-controller` - **FIXED**: Updates handleAccountRemoved to update the selected account to the most recent one, if the account being removed is currently the selected account, all within the same step. ## Checklist - [x] I've updated the test suite for new or updated code as appropriate - [x] I've updated documentation (JSDoc, Markdown, etc.) for new or updated code as appropriate - [ ] I've highlighted breaking changes using the "BREAKING" category above as appropriate --------- Co-authored-by: Charly Chevalier <charly.chevalier@consensys.net>
diff --git a/packages/accounts-controller/src/AccountsController.test.ts b/packages/accounts-controller/src/AccountsController.test.ts
@@ -1078,16 +1078,10 @@ describe('AccountsController', () => {
           initialState: {
             internalAccounts: {
               accounts: {
-                'missing-account': {
-                  id: 'missing-account',
+                'missing-account': createMockInternalAccount({
                   address: '0x999',
-                  metadata: {
-                    keyring: {
-                      type: KeyringTypes.hd,
-                    },
-                  },
-                  [mockAccount2.id]: mockAccount2WithoutLastSelected,
-                } as unknown as InternalAccount,
+                  id: 'missing-account',
+                }),
                 [mockAccount.id]: mockAccountWithoutLastSelected,
                 [mockAccount2.id]: mockAccount2WithoutLastSelected,
               },
@@ -1850,6 +1844,32 @@ describe('AccountsController', () => {
         'No EVM accounts',
       );
     });
+
+    it('handle the edge case of undefined accountId during onboarding', async () => {
+      const { accountsController } = setupAccountsController({
+        initialState: {
+          internalAccounts: {
+            accounts: {},
+            selectedAccount: '',
+          },
+        },
+      });
+
+      expect(accountsController.getSelectedAccount()).toStrictEqual({
+        id: '',
+        address: '',
+        options: {},
+        methods: [],
+        type: EthAccountType.Eoa,
+        metadata: {
+          name: '',
+          keyring: {
+            type: '',
+          },
+          importTime: 0,
+        },
+      });
+    });
   });
 
   describe('getSelectedMultichainAccount', () => {
@@ -2060,33 +2080,6 @@ describe('AccountsController', () => {
         `Account Id "${accountId}" not found`,
       );
     });
-
-    it('handle the edge case of undefined accountId during onboarding', async () => {
-      const { accountsController } = setupAccountsController({
-        initialState: {
-          internalAccounts: {
-            accounts: { [mockAccount.id]: mockAccount },
-            selectedAccount: mockAccount.id,
-          },
-        },
-      });
-
-      // @ts-expect-error forcing undefined accountId
-      expect(accountsController.getAccountExpect(undefined)).toStrictEqual({
-        id: '',
-        address: '',
-        options: {},
-        methods: [],
-        type: EthAccountType.Eoa,
-        metadata: {
-          name: '',
-          keyring: {
-            type: '',
-          },
-          importTime: 0,
-        },
-      });
-    });
   });
 
   describe('setSelectedAccount', () => {
diff --git a/packages/accounts-controller/src/AccountsController.ts b/packages/accounts-controller/src/AccountsController.ts
@@ -35,7 +35,6 @@ import {
 import type { Draft } from 'immer';
 
 import {
-  deepCloneDraft,
   getUUIDFromAddressOfNormalAccount,
   isNormalKeyringType,
   keyringTypeToName,
@@ -271,9 +270,22 @@ export class AccountsController extends BaseController<
    * @throws An error if the account ID is not found.
    */
   getAccountExpect(accountId: string): InternalAccount {
+    const account = this.getAccount(accountId);
+    if (account === undefined) {
+      throw new Error(`Account Id "${accountId}" not found`);
+    }
+    return account;
+  }
+
+  /**
+   * Returns the last selected EVM account.
+   *
+   * @returns The selected internal account.
+   */
+  getSelectedAccount(): InternalAccount {
     // Edge case where the extension is setup but the srp is not yet created
     // certain ui elements will query the selected address before any accounts are created.
-    if (!accountId) {
+    if (this.state.internalAccounts.selectedAccount === '') {
       return {
         id: '',
         address: '',
@@ -290,19 +302,6 @@ export class AccountsController extends BaseController<
       };
     }
 
-    const account = this.getAccount(accountId);
-    if (account === undefined) {
-      throw new Error(`Account Id "${accountId}" not found`);
-    }
-    return account;
-  }
-
-  /**
-   * Returns the last selected evm account.
-   *
-   * @returns The selected internal account.
-   */
-  getSelectedAccount(): InternalAccount {
     const selectedAccount = this.getAccountExpect(
       this.state.internalAccounts.selectedAccount,
     );
@@ -412,12 +411,7 @@ export class AccountsController extends BaseController<
         ...account,
         metadata: { ...account.metadata, name: accountName },
       };
-      // FIXME: deep clone of old state to get around Type instantiation is excessively deep and possibly infinite.
-      const newState = deepCloneDraft(currentState);
-
-      newState.internalAccounts.accounts[accountId] = internalAccount;
-
-      return newState;
+      currentState.internalAccounts.accounts[accountId] = internalAccount;
     });
   }
 
@@ -474,12 +468,7 @@ export class AccountsController extends BaseController<
     }, {} as Record<string, InternalAccount>);
 
     this.update((currentState: Draft<AccountsControllerState>) => {
-      // FIXME: deep clone of old state to get around Type instantiation is excessively deep and possibly infinite.
-      const newState = deepCloneDraft(currentState);
-
-      newState.internalAccounts.accounts = accounts;
-
-      return newState;
+      currentState.internalAccounts.accounts = accounts;
     });
   }
 
@@ -491,12 +480,7 @@ export class AccountsController extends BaseController<
   loadBackup(backup: AccountsControllerState): void {
     if (backup.internalAccounts) {
       this.update((currentState: Draft<AccountsControllerState>) => {
-        // FIXME: deep clone of old state to get around Type instantiation is excessively deep and possibly infinite.
-        const newState = deepCloneDraft(currentState);
-
-        newState.internalAccounts = backup.internalAccounts;
-
-        return newState;
+        currentState.internalAccounts = backup.internalAccounts;
       });
     }
   }
@@ -717,41 +701,58 @@ export class AccountsController extends BaseController<
         }
       }
 
-      if (deletedAccounts.length > 0) {
-        for (const account of deletedAccounts) {
-          this.#handleAccountRemoved(account.id);
+      this.update((currentState: Draft<AccountsControllerState>) => {
+        if (deletedAccounts.length > 0) {
+          for (const account of deletedAccounts) {
+            currentState.internalAccounts.accounts = this.#handleAccountRemoved(
+              currentState.internalAccounts.accounts,
+              account.id,
+            );
+          }
         }
-      }
 
-      if (addedAccounts.length > 0) {
-        for (const account of addedAccounts) {
-          this.#handleNewAccountAdded(account);
+        if (addedAccounts.length > 0) {
+          for (const account of addedAccounts) {
+            currentState.internalAccounts.accounts =
+              this.#handleNewAccountAdded(
+                currentState.internalAccounts.accounts,
+                account,
+              );
+          }
         }
-      }
 
-      // handle if the selected account was deleted
-      if (!this.getAccount(this.state.internalAccounts.selectedAccount)) {
-        const [accountToSelect] = this.listAccounts().sort(
-          (accountA, accountB) => {
-            // sort by lastSelected descending
-            return (
-              (accountB.metadata.lastSelected ?? 0) -
-              (accountA.metadata.lastSelected ?? 0)
-            );
-          },
+        // We don't use list accounts because it is not the updated state yet.
+        const existingAccounts = Object.values(
+          currentState.internalAccounts.accounts,
         );
 
-        // if the accountToSelect is undefined, then there are no accounts
-        // it mean the keyring was reinitialized.
-        if (!accountToSelect) {
-          this.update((currentState: Draft<AccountsControllerState>) => {
+        // handle if the selected account was deleted
+        if (
+          !currentState.internalAccounts.accounts[
+            this.state.internalAccounts.selectedAccount
+          ]
+        ) {
+          // if currently selected account is undefined and there are no accounts
+          // it mean the keyring was reinitialized.
+          if (existingAccounts.length === 0) {
             currentState.internalAccounts.selectedAccount = '';
-          });
-          return;
-        }
+            return;
+          }
 
-        this.setSelectedAccount(accountToSelect.id);
-      }
+          // at this point, we know that `existingAccounts.length` is > 0, so
+          // `accountToSelect` won't be `undefined`!
+          const [accountToSelect] = existingAccounts.sort(
+            (accountA, accountB) => {
+              // sort by lastSelected descending
+              return (
+                (accountB.metadata.lastSelected ?? 0) -
+                (accountA.metadata.lastSelected ?? 0)
+              );
+            },
+          );
+          currentState.internalAccounts.selectedAccount = accountToSelect.id;
+        }
+      });
     }
   }
 
@@ -786,10 +787,11 @@ export class AccountsController extends BaseController<
   /**
    * Returns the list of accounts for a given keyring type.
    * @param keyringType - The type of keyring.
+   * @param accounts - Accounts to filter by keyring type.
    * @returns The list of accounts associcated with this keyring type.
    */
-  #getAccountsByKeyringType(keyringType: string) {
-    return this.listAccounts().filter((internalAccount) => {
+  #getAccountsByKeyringType(keyringType: string, accounts?: InternalAccount[]) {
+    return (accounts ?? this.listAccounts()).filter((internalAccount) => {
       // We do consider `hd` and `simple` keyrings to be of same type. So we check those 2 types
       // to group those accounts together!
       if (
@@ -833,11 +835,18 @@ export class AccountsController extends BaseController<
   /**
    * Returns the next account number for a given keyring type.
    * @param keyringType - The type of keyring.
+   * @param accounts - Existing accounts to check for the next available account number.
    * @returns An object containing the account prefix and index to use.
    */
-  getNextAvailableAccountName(keyringType: string = KeyringTypes.hd): string {
+  getNextAvailableAccountName(
+    keyringType: string = KeyringTypes.hd,
+    accounts?: InternalAccount[],
+  ): string {
     const keyringName = keyringTypeToName(keyringType);
-    const keyringAccounts = this.#getAccountsByKeyringType(keyringType);
+    const keyringAccounts = this.#getAccountsByKeyringType(
+      keyringType,
+      accounts,
+    );
     const lastDefaultIndexUsedForKeyringType = keyringAccounts.reduce(
       (maxInternalAccountIndex, internalAccount) => {
         // We **DO NOT USE** `\d+` here to only consider valid "human"
@@ -888,9 +897,14 @@ export class AccountsController extends BaseController<
    * Handles the addition of a new account to the controller.
    * If the account is not a Snap Keyring account, generates an internal account for it and adds it to the controller.
    * If the account is a Snap Keyring account, retrieves the account from the keyring and adds it to the controller.
+   * @param accountsState - AccountsController accounts state that is to be mutated.
    * @param account - The address and keyring type object of the new account.
+   * @returns The updated AccountsController accounts state.
    */
-  #handleNewAccountAdded(account: AddressAndKeyringTypeObject) {
+  #handleNewAccountAdded(
+    accountsState: AccountsControllerState['internalAccounts']['accounts'],
+    account: AddressAndKeyringTypeObject,
+  ): AccountsControllerState['internalAccounts']['accounts'] {
     let newAccount: InternalAccount;
     if (account.type !== KeyringTypes.snap) {
       newAccount = this.#generateInternalAccountForNonSnapAccount(
@@ -909,41 +923,41 @@ export class AccountsController extends BaseController<
 
       // The snap deleted the account before the keyring controller could add it
       if (!newAccount) {
-        return;
+        return accountsState;
       }
     }
 
     // Get next account name available for this given keyring
     const accountName = this.getNextAvailableAccountName(
       newAccount.metadata.keyring.type,
+      Object.values(accountsState),
     );
 
-    this.update((currentState: Draft<AccountsControllerState>) => {
-      // FIXME: deep clone of old state to get around Type instantiation is excessively deep and possibly infinite.
-      const newState = deepCloneDraft(currentState);
-
-      newState.internalAccounts.accounts[newAccount.id] = {
-        ...newAccount,
-        metadata: {
-          ...newAccount.metadata,
-          name: accountName,
-          importTime: Date.now(),
-          lastSelected: 0,
-        },
-      };
+    accountsState[newAccount.id] = {
+      ...newAccount,
+      metadata: {
+        ...newAccount.metadata,
+        name: accountName,
+        importTime: Date.now(),
+        lastSelected: 0,
+      },
+    };
 
-      return newState;
-    });
+    return accountsState;
   }
 
   /**
    * Handles the removal of an account from the internal accounts list.
+   * @param accountsState - AccountsController accounts state that is to be mutated.
    * @param accountId - The ID of the account to be removed.
+   * @returns The updated AccountsController state.
    */
-  #handleAccountRemoved(accountId: string) {
-    this.update((currentState: Draft<AccountsControllerState>) => {
-      delete currentState.internalAccounts.accounts[accountId];
-    });
+  #handleAccountRemoved(
+    accountsState: AccountsControllerState['internalAccounts']['accounts'],
+    accountId: string,
+  ): AccountsControllerState['internalAccounts']['accounts'] {
+    delete accountsState[accountId];
+    return accountsState;
   }
 
   /**
diff --git a/packages/accounts-controller/src/utils.ts b/packages/accounts-controller/src/utils.ts
