feat(account-group): add groupByWalletType rule and listAccountGroups function

Author: PatrykLucka

packages/account-group-controller/src/AccountGroupController.test.ts

@@ -19,7 +19,10 @@ import {
   type AllowedActions,
   type AllowedEvents,
 } from './AccountGroupController';
-import type { AccountGroupControllerMessenger } from './AccountGroupController';
+import type {
+  AccountGroupControllerMessenger,
+  AccountGroupId,
+} from './AccountGroupController';
 
 const ETH_EOA_METHODS = [
   EthMethod.PersonalSign,
@@ -31,7 +34,9 @@ const ETH_EOA_METHODS = [
 ] as const;
 
 /**
+ * Creates a new root messenger instance for testing.
  *
+ * @returns A new Messenger instance.
  */
 function getRootMessenger() {
   return new Messenger<
@@ -41,8 +46,10 @@ function getRootMessenger() {
 }
 
 /**
+ * Retrieves a restricted messenger for the AccountGroupController.
  *
- * @param messenger
+ * @param messenger - The root messenger instance. Defaults to a new Messenger created by getRootMessenger().
+ * @returns The restricted messenger for the AccountGroupController.
  */
 function getAccountGroupControllerMessenger(
   messenger = getRootMessenger(),
@@ -58,11 +65,12 @@ function getAccountGroupControllerMessenger(
 }
 
 /**
+ * Sets up the AccountGroupController for testing.
  *
- * @param options0
- * @param options0.initialState
- * @param options0.messenger
- * @param options0.state
+ * @param options - Configuration options for setup.
+ * @param options.state - Partial initial state for the controller. Defaults to empty object.
+ * @param options.messenger - An optional messenger instance to use. Defaults to a new Messenger.
+ * @returns An object containing the controller instance and the messenger.
  */
 function setup({
   state = {},
@@ -168,39 +176,184 @@ const MOCK_SNAP_ACCOUNT_2: InternalAccount = {
   },
 };
 
+const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = {
+  id: 'mock-hardware-id-1',
+  address: '0x123',
+  options: {},
+  methods: [...ETH_EOA_METHODS],
+  type: EthAccountType.Eoa,
+  scopes: [EthScope.Eoa],
+  metadata: {
+    name: '',
+    keyring: { type: KeyringTypes.ledger },
+    importTime: 1691565967600,
+    lastSelected: 1955565967656,
+  },
+};
+
 describe('AccountGroupController', () => {
-  it('group accounts according to the rules', async () => {
-    const { controller, messenger } = setup({});
-
-    messenger.registerActionHandler(
-      'AccountsController:listMultichainAccounts',
-      () => {
-        return [
-          MOCK_HD_ACCOUNT_1,
-          MOCK_HD_ACCOUNT_2,
-          MOCK_SNAP_ACCOUNT_1,
-          MOCK_SNAP_ACCOUNT_2,
-        ];
-      },
-    );
-
-    await controller.init();
-
-    expect(controller.state).toStrictEqual({
-      accountGroups: {
-        groups: {
-          'mock-keyring-id-1': {
-            [DEFAULT_SUB_GROUP]: [MOCK_HD_ACCOUNT_1.id],
+  describe('updateAccountGroups', () => {
+    it('should group accounts by entropy source, then snapId, then wallet type', async () => {
+      const { controller, messenger } = setup({});
+
+      messenger.registerActionHandler(
+        'AccountsController:listMultichainAccounts',
+        () => {
+          return [
+            MOCK_HD_ACCOUNT_1,
+            MOCK_HD_ACCOUNT_2,
+            MOCK_SNAP_ACCOUNT_1,
+            MOCK_SNAP_ACCOUNT_2,
+            MOCK_HARDWARE_ACCOUNT_1,
+          ];
+        },
+      );
+
+      await controller.updateAccountGroups();
+
+      expect(controller.state.accountGroups.groups).toStrictEqual({
+        'mock-keyring-id-1': {
+          [DEFAULT_SUB_GROUP]: [MOCK_HD_ACCOUNT_1.id],
+        },
+        'mock-keyring-id-2': {
+          [DEFAULT_SUB_GROUP]: [MOCK_HD_ACCOUNT_2.id, MOCK_SNAP_ACCOUNT_1.id],
+        },
+        'mock-snap-id-2': {
+          [DEFAULT_SUB_GROUP]: [MOCK_SNAP_ACCOUNT_2.id],
+        },
+        [KeyringTypes.ledger]: {
+          [DEFAULT_SUB_GROUP]: [MOCK_HARDWARE_ACCOUNT_1.id],
+        },
+      });
+    });
+
+    it('should warn and fall back to wallet type grouping if an HD account is missing entropySource', async () => {
+      const consoleWarnSpy = jest
+        .spyOn(console, 'warn')
+        .mockImplementation(() => undefined);
+      const { controller, messenger } = setup({});
+
+      const mockHdAccountWithoutEntropy: InternalAccount = {
+        id: 'hd-account-no-entropy',
+        address: '0xHDADD',
+        metadata: {
+          name: 'HD Account Without Entropy',
+          keyring: {
+            type: KeyringTypes.hd,
+          },
+          importTime: Date.now(),
+          lastSelected: Date.now(),
+        },
+        methods: [...ETH_EOA_METHODS],
+        options: {},
+        type: EthAccountType.Eoa,
+        scopes: [EthScope.Eoa],
+      };
+
+      const listAccountsMock = jest
+        .fn()
+        .mockReturnValue([mockHdAccountWithoutEntropy]);
+      messenger.registerActionHandler(
+        'AccountsController:listMultichainAccounts',
+        listAccountsMock,
+      );
+
+      await controller.updateAccountGroups();
+
+      expect(consoleWarnSpy).toHaveBeenCalledWith(
+        "! Found an HD account with no entropy source: account won't be associated to its wallet",
+      );
+
+      expect(
+        controller.state.accountGroups.groups[KeyringTypes.hd]?.[
+          DEFAULT_SUB_GROUP
+        ],
+      ).toContain(mockHdAccountWithoutEntropy.id);
+
+      expect(
+        controller.state.accountGroups.groups[
+          undefined as unknown as AccountGroupId
+        ],
+      ).toBeUndefined();
+
+      consoleWarnSpy.mockRestore();
+    });
+  });
+
+  describe('listAccountGroups', () => {
+    it('should return an empty array if no groups exist', async () => {
+      const { controller } = setup({
+        state: {
+          accountGroups: { groups: {} },
+          accountGroupsMetadata: {},
+        },
+      });
+
+      const result = await controller.listAccountGroups();
+      expect(result).toEqual([]);
+    });
+
+    it('should correctly map group data and metadata to AccountGroup objects', async () => {
+      const group1Id = 'group-id-1' as AccountGroupId;
+      const group2Id = 'group-id-2' as AccountGroupId;
+
+      const initialState: AccountGroupControllerState = {
+        accountGroups: {
+          groups: {
+            [group1Id]: {
+              [DEFAULT_SUB_GROUP]: ['account-1', 'account-2'],
+            },
+            [group2Id]: {
+              'sub-group-x': ['account-3'],
+            },
           },
-          'mock-keyring-id-2': {
-            [DEFAULT_SUB_GROUP]: [MOCK_HD_ACCOUNT_2.id, MOCK_SNAP_ACCOUNT_1.id],
+        },
+        accountGroupsMetadata: {
+          [group1Id]: { name: 'Group 1 Name' },
+          [group2Id]: { name: 'Group 2 Name' },
+        },
+      };
+
+      const { controller } = setup({ state: initialState });
+      const result = await controller.listAccountGroups();
+
+      expect(result).toEqual([
+        {
+          id: group1Id,
+          name: 'Group 1 Name',
+          subGroups: {
+            [DEFAULT_SUB_GROUP]: ['account-1', 'account-2'],
           },
-          'mock-snap-id-2': {
-            [DEFAULT_SUB_GROUP]: [MOCK_SNAP_ACCOUNT_2.id],
+        },
+        {
+          id: group2Id,
+          name: 'Group 2 Name',
+          subGroups: {
+            'sub-group-x': ['account-3'],
           },
         },
-        metadata: {},
-      },
-    } as AccountGroupControllerState);
+      ]);
+    });
+
+    it('should throw a TypeError if metadata is missing for a group', async () => {
+      const groupIdWithMissingMetadata = 'group-missing-meta' as AccountGroupId;
+      const initialState: Partial<AccountGroupControllerState> = {
+        accountGroups: {
+          groups: {
+            [groupIdWithMissingMetadata]: {
+              [DEFAULT_SUB_GROUP]: ['account-x'],
+            },
+          },
+        },
+        // Metadata for groupIdWithMissingMetadata is deliberately omitted
+        accountGroupsMetadata: {},
+      };
+
+      const { controller } = setup({ state: initialState });
+
+      // Current implementation will throw: Cannot read properties of undefined (reading 'name')
+      // which is a TypeError.
+      await expect(controller.listAccountGroups()).rejects.toThrow(TypeError);
+    });
   });
 });

packages/account-group-controller/src/AccountGroupController.ts

@@ -21,6 +21,14 @@ export type AccountGroupId = string;
 // NOTES:
 // - Maybe add a `metadata` / `flags` for each groups (or at least, top-level ones)
 
+export type AccountGroup = {
+  id: AccountGroupId;
+  name: string;
+  subGroups: {
+    [accountSubGroup: AccountGroupId]: AccountId[];
+  };
+};
+
 export type AccountGroupMetadata = {
   name: string;
 };
@@ -34,10 +42,9 @@ export type AccountGroupControllerState = {
         [accountSubGroup: AccountGroupId]: AccountId[]; // Blockchain Accounts
       };
     };
-    // QUESTION: Should we have some metadata here indexed by `AccountGroupId` or should we
-    // group them on `groups` and sub-groups? The reasoning of having them separate would be
-    // to be able to persist them and just re-create the groups at runtime.
-    metadata: Record<AccountGroupId, AccountGroupMetadata>;
+  };
+  accountGroupsMetadata: {
+    [accountGroup: AccountGroupId]: AccountGroupMetadata;
   };
 };
 
@@ -46,9 +53,15 @@ export type AccountGroupControllerGetStateAction = ControllerGetStateAction<
   AccountGroupControllerState
 >;
 
+export type AccountGroupControllerListAccountGroupsAction = {
+  type: `${typeof controllerName}:listAccountGroups`;
+  handler: AccountGroupController['listAccountGroups'];
+};
+
 export type AllowedActions = AccountsControllerListMultichainAccountsAction;
 
-export type AccountGroupControllerActions = never;
+export type AccountGroupControllerActions =
+  AccountGroupControllerListAccountGroupsAction;
 
 export type AccountGroupControllerChangeEvent = ControllerStateChangeEvent<
   typeof controllerName,
@@ -75,6 +88,10 @@ const accountGroupControllerMetadata: StateMetadata<AccountGroupControllerState>
       persist: false, // We do re-recompute this state everytime.
       anonymous: false,
     },
+    accountGroupsMetadata: {
+      persist: false, // TODO: Change it to true once we have data to persist.
+      anonymous: false,
+    },
   };
 
 /**
@@ -86,8 +103,8 @@ export function getDefaultAccountGroupControllerState(): AccountGroupControllerS
   return {
     accountGroups: {
       groups: {},
-      metadata: {},
     },
+    accountGroupsMetadata: {},
   };
 }
 
@@ -183,12 +200,18 @@ export class AccountGroupController extends BaseController<
     return undefined;
   }
 
-  async init(): Promise<void> {
+  #groupByWalletType(account: InternalAccount): AccountGroupId | undefined {
+    return account.metadata.keyring.type as AccountGroupId;
+  }
+
+  async updateAccountGroups(): Promise<void> {
     const rules = [
       // 1. We group by entropy-source
       (account: InternalAccount) => this.#groupByEntropySource(account),
       // 2. We group by Snap ID
       (account: InternalAccount) => this.#groupBySnapId(account),
+      // 3. We group by wallet type
+      (account: InternalAccount) => this.#groupByWalletType(account),
     ];
     const groups: AccountGroupControllerState['accountGroups']['groups'] = {};
 
@@ -219,6 +242,19 @@ export class AccountGroupController extends BaseController<
     });
   }
 
+  async listAccountGroups(): Promise<AccountGroup[]> {
+    return Object.keys(this.state.accountGroups.groups).map(
+      (groupId: AccountGroupId) => {
+        const subGroups = this.state.accountGroups.groups[groupId];
+        return {
+          id: groupId,
+          name: this.state.accountGroupsMetadata[groupId].name,
+          subGroups,
+        };
+      },
+    );
+  }
+
   /**
    * Lists the multichain accounts coming from the `AccountsController`.
    *