Make addNewAccount idempotent (#1298)

* feat: addNewAccount idempotent * test: add assertion for test case * refactor: apply review suggestions * test: add coverage for call without accountCount
MetaMask · Oct 11, 2023 · fed201e · fed201e
1 parent 9bc12b8
commit fed201e
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 20 deletions.
diff --git a/packages/keyring-controller/src/KeyringController.test.ts b/packages/keyring-controller/src/KeyringController.test.ts
@@ -91,26 +91,73 @@ describe('KeyringController', () => {
   });
 
   describe('addNewAccount', () => {
-    it('should add new account', async () => {
-      const { keyringState: currentKeyringMemState, addedAccountAddress } =
-        await keyringController.addNewAccount();
-      expect(initialState.keyrings).toHaveLength(1);
-      expect(initialState.keyrings[0].accounts).not.toStrictEqual(
-        currentKeyringMemState.keyrings[0].accounts,
-      );
-      expect(currentKeyringMemState.keyrings[0].accounts).toHaveLength(2);
-      expect(initialState.keyrings[0].accounts).not.toContain(
-        addedAccountAddress,
-      );
-      expect(addedAccountAddress).toBe(
-        currentKeyringMemState.keyrings[0].accounts[1],
-      );
-      expect(
-        preferences.updateIdentities.calledWith(
+    describe('when accountCount is not provided', () => {
+      it('should add new account', async () => {
+        const { keyringState: currentKeyringMemState, addedAccountAddress } =
+          await keyringController.addNewAccount();
+        expect(initialState.keyrings).toHaveLength(1);
+        expect(initialState.keyrings[0].accounts).not.toStrictEqual(
           currentKeyringMemState.keyrings[0].accounts,
-        ),
-      ).toBe(true);
-      expect(preferences.setSelectedAddress.called).toBe(false);
+        );
+        expect(currentKeyringMemState.keyrings[0].accounts).toHaveLength(2);
+        expect(initialState.keyrings[0].accounts).not.toContain(
+          addedAccountAddress,
+        );
+        expect(addedAccountAddress).toBe(
+          currentKeyringMemState.keyrings[0].accounts[1],
+        );
+        expect(
+          preferences.updateIdentities.calledWith(
+            currentKeyringMemState.keyrings[0].accounts,
+          ),
+        ).toBe(true);
+        expect(preferences.setSelectedAddress.called).toBe(false);
+      });
+    });
+
+    describe('when accountCount is provided', () => {
+      it('should add new account if accountCount is in sequence', async () => {
+        const { keyringState: currentKeyringMemState, addedAccountAddress } =
+          await keyringController.addNewAccount(
+            initialState.keyrings[0].accounts.length,
+          );
+        expect(initialState.keyrings).toHaveLength(1);
+        expect(initialState.keyrings[0].accounts).not.toStrictEqual(
+          currentKeyringMemState.keyrings[0].accounts,
+        );
+        expect(currentKeyringMemState.keyrings[0].accounts).toHaveLength(2);
+        expect(initialState.keyrings[0].accounts).not.toContain(
+          addedAccountAddress,
+        );
+        expect(addedAccountAddress).toBe(
+          currentKeyringMemState.keyrings[0].accounts[1],
+        );
+        expect(
+          preferences.updateIdentities.calledWith(
+            currentKeyringMemState.keyrings[0].accounts,
+          ),
+        ).toBe(true);
+        expect(preferences.setSelectedAddress.called).toBe(false);
+      });
+
+      it('should throw an error if passed accountCount param is out of sequence', async () => {
+        const accountCount = initialState.keyrings[0].accounts.length;
+        await expect(
+          keyringController.addNewAccount(accountCount + 1),
+        ).rejects.toThrow('Account out of sequence');
+      });
+
+      it('should not add a new account if called twice with the same accountCount param', async () => {
+        const accountCount = initialState.keyrings[0].accounts.length;
+        const { addedAccountAddress: firstAccountAdded } =
+          await keyringController.addNewAccount(accountCount);
+        const { keyringState, addedAccountAddress: secondAccountAdded } =
+          await keyringController.addNewAccount(accountCount);
+        expect(firstAccountAdded).toBe(secondAccountAdded);
+        expect(keyringState.keyrings[0].accounts).toHaveLength(
+          accountCount + 1,
+        );
+      });
     });
   });
 

diff --git a/packages/keyring-controller/src/KeyringController.ts b/packages/keyring-controller/src/KeyringController.ts
@@ -205,10 +205,12 @@ export class KeyringController extends BaseController<
   /**
    * Adds a new account to the default (first) HD seed phrase keyring.
    *
+   * @param accountCount - Number of accounts before adding a new one, used to
+   * make the method idempotent.
    * @returns Promise resolving to keyring current state and added account
    * address.
    */
-  async addNewAccount(): Promise<{
+  async addNewAccount(accountCount?: number): Promise<{
     keyringState: KeyringMemState;
     addedAccountAddress: string;
   }> {
@@ -218,6 +220,19 @@ export class KeyringController extends BaseController<
       throw new Error('No HD keyring found');
     }
     const oldAccounts = await this.#keyring.getAccounts();
+
+    if (accountCount && oldAccounts.length !== accountCount) {
+      if (accountCount > oldAccounts.length) {
+        throw new Error('Account out of sequence');
+      }
+      // we return the account already existing at index `accountCount`
+      const primaryKeyringAccounts = await primaryKeyring.getAccounts();
+      return {
+        keyringState: await this.fullUpdate(),
+        addedAccountAddress: primaryKeyringAccounts[accountCount],
+      };
+    }
+
     await this.#keyring.addNewAccount(primaryKeyring);
     const newAccounts = await this.#keyring.getAccounts();