fix: Tron addresses not being generated for all BIP44 accounts

Author: ulissesferreira

packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts

@@ -33,51 +33,40 @@ class MockTronKeyring {
     this.accounts = accounts;
   }
 
-  #getIndexFromDerivationPath(derivationPath: string): number {
-    // eslint-disable-next-line prefer-regex-literals
-    const derivationPathIndexRegex = new RegExp(
-      "^m/44'/195'/0'/(?<index>[0-9]+)'$",
-      'u',
-    );
-
-    const matched = derivationPath.match(derivationPathIndexRegex);
-    if (matched?.groups?.index === undefined) {
-      throw new Error('Unable to extract index');
-    }
-
-    const { index } = matched.groups;
-    return Number(index);
-  }
-
   createAccount: SnapKeyring['createAccount'] = jest
     .fn()
-    .mockImplementation((_, { derivationPath }) => {
-      if (derivationPath !== undefined) {
-        const index = this.#getIndexFromDerivationPath(derivationPath);
-        const found = this.accounts.find(
-          (account) =>
-            isBip44Account(account) &&
-            account.options.entropy.groupIndex === index,
-        );
-
-        if (found) {
-          return found; // Idempotent.
-        }
+    .mockImplementation((_, { index }) => {
+      // Use the provided index or fallback to accounts length
+      const groupIndex = index !== undefined ? index : this.accounts.length;
+
+      // Check if an account already exists for this group index (idempotent behavior)
+      const found = this.accounts.find(
+        (account) =>
+          isBip44Account(account) &&
+          account.options.entropy.groupIndex === groupIndex,
+      );
+
+      if (found) {
+        return found; // Idempotent.
       }
 
+      // Create new account with the correct group index
       const account = MockAccountBuilder.from(MOCK_TRX_ACCOUNT_1)
         .withUuid()
         .withAddressSuffix(`${this.accounts.length}`)
-        .withGroupIndex(this.accounts.length)
+        .withGroupIndex(groupIndex)
         .get();
       this.accounts.push(account);
 
       return account;
     });
+
+  // Add discoverAccounts method to match the provider's usage
+  discoverAccounts = jest.fn().mockResolvedValue([]);
 }
 
 /**
- * Sets up a SolAccountProvider for testing.
+ * Sets up a TrxAccountProvider for testing.
  *
  * @param options - Configuration options for setup.
  * @param options.messenger - An optional messenger instance to use. Defaults to a new Messenger.
@@ -98,6 +87,7 @@ function setup({
     handleRequest: jest.Mock;
     keyring: {
       createAccount: jest.Mock;
+      discoverAccounts: jest.Mock;
     };
   };
 } {
@@ -108,11 +98,17 @@ function setup({
     () => accounts,
   );
 
-  const mockHandleRequest = jest
-    .fn()
-    .mockImplementation((address: string) =>
-      keyring.accounts.find((account) => account.address === address),
+  const mockHandleRequest = jest.fn().mockImplementation((request) => {
+    // Handle KeyringClient discoverAccounts calls
+    if (request.request?.method === 'keyring_discoverAccounts') {
+      // Return the keyring's discoverAccounts result directly
+      return keyring.discoverAccounts();
+    }
+    // Handle other requests (fallback for legacy compatibility)
+    return keyring.accounts.find(
+      (account) => account.address === request.address,
     );
+  });
   messenger.registerActionHandler(
     'SnapController:handleRequest',
     mockHandleRequest,
@@ -143,6 +139,7 @@ function setup({
       handleRequest: mockHandleRequest,
       keyring: {
         createAccount: keyring.createAccount as jest.Mock,
+        discoverAccounts: keyring.discoverAccounts as jest.Mock,
       },
     },
   };
@@ -235,7 +232,7 @@ describe('TrxAccountProvider', () => {
   });
 
   // Skip this test for now, since we manually inject those options upon
-  // account creation, so it cannot fails (until the Solana Snap starts
+  // account creation, so it cannot fails (until the Tron Snap starts
   // using the new typed options).
   // eslint-disable-next-line jest/no-disabled-tests
   it.skip('throws if the created account is not BIP-44 compatible', async () => {
@@ -263,7 +260,9 @@ describe('TrxAccountProvider', () => {
     });
 
     // Simulate one discovered account at the requested index.
-    mocks.handleRequest.mockReturnValue([MOCK_TRX_DISCOVERED_ACCOUNT_1]);
+    mocks.keyring.discoverAccounts.mockResolvedValue([
+      MOCK_TRX_DISCOVERED_ACCOUNT_1,
+    ]);
 
     const discovered = await provider.discoverAccounts({
       entropySource: MOCK_HD_KEYRING_1.metadata.id,
@@ -283,7 +282,9 @@ describe('TrxAccountProvider', () => {
     });
 
     // Simulate one discovered account — should resolve to the existing one
-    mocks.handleRequest.mockReturnValue([MOCK_TRX_DISCOVERED_ACCOUNT_1]);
+    mocks.keyring.discoverAccounts.mockResolvedValue([
+      MOCK_TRX_DISCOVERED_ACCOUNT_1,
+    ]);
 
     const discovered = await provider.discoverAccounts({
       entropySource: MOCK_HD_KEYRING_1.metadata.id,
@@ -298,7 +299,7 @@ describe('TrxAccountProvider', () => {
       accounts: [],
     });
 
-    mocks.handleRequest.mockReturnValue([]);
+    mocks.keyring.discoverAccounts.mockResolvedValue([]);
 
     const discovered = await provider.discoverAccounts({
       entropySource: MOCK_HD_KEYRING_1.metadata.id,

packages/multichain-account-service/src/providers/TrxAccountProvider.ts

@@ -1,7 +1,7 @@
 import { assertIsBip44Account, type Bip44Account } from '@metamask/account-api';
 import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api';
 import { TrxAccountType, TrxScope } from '@metamask/keyring-api';
-import { KeyringAccountEntropyTypeOption } from '@metamask/keyring-api';
+// import { KeyringAccountEntropyTypeOption } from '@metamask/keyring-api';
 import { KeyringTypes } from '@metamask/keyring-controller';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 import { KeyringClient } from '@metamask/keyring-snap-client';
@@ -83,47 +83,26 @@ export class TrxAccountProvider extends SnapAccountProvider {
     );
   }
 
-  async #createAccount({
+  async createAccounts({
     entropySource,
-    groupIndex,
-    derivationPath,
+    groupIndex: index,
   }: {
     entropySource: EntropySourceId;
     groupIndex: number;
-    derivationPath: string;
-  }): Promise<Bip44Account<KeyringAccount>> {
+  }): Promise<Bip44Account<KeyringAccount>[]> {
     const createAccount = await this.getRestrictedSnapAccountCreator();
+
     const account = await withTimeout(
-      createAccount({ entropySource, derivationPath }),
+      createAccount({
+        entropySource,
+        index,
+        addressType: TrxAccountType.Eoa,
+        scope: TrxScope.Mainnet,
+      }),
       this.#config.createAccounts.timeoutMs,
     );
 
-    // Ensure entropy is present before type assertion validation
-    account.options.entropy = {
-      type: KeyringAccountEntropyTypeOption.Mnemonic,
-      id: entropySource,
-      groupIndex,
-      derivationPath,
-    };
-
     assertIsBip44Account(account);
-    return account;
-  }
-
-  async createAccounts({
-    entropySource,
-    groupIndex,
-  }: {
-    entropySource: EntropySourceId;
-    groupIndex: number;
-  }): Promise<Bip44Account<KeyringAccount>[]> {
-    const derivationPath = `m/44'/195'/0'/${groupIndex}'`;
-    const account = await this.#createAccount({
-      entropySource,
-      groupIndex,
-      derivationPath,
-    });
-
     return [account];
   }
 
@@ -154,15 +133,10 @@ export class TrxAccountProvider extends SnapAccountProvider {
       return [];
     }
 
-    const createdAccounts = await Promise.all(
-      discoveredAccounts.map((d) =>
-        this.#createAccount({
-          entropySource,
-          groupIndex,
-          derivationPath: d.derivationPath,
-        }),
-      ),
-    );
+    const createdAccounts = await this.createAccounts({
+      entropySource,
+      groupIndex,
+    });
 
     return createdAccounts;
   }