feat: support multiple accounts in SNAP

packages/starknet-snap/src/index.tsx

@@ -18,7 +18,6 @@ import { getErc20TokenBalance } from './getErc20TokenBalance';
 import { getStarkName } from './getStarkName';
 import { getStoredErc20Tokens } from './getStoredErc20Tokens';
 import { getStoredNetworks } from './getStoredNetworks';
-import { getStoredUserAccounts } from './getStoredUserAccounts';
 import { getValue } from './getValue';
 import { homePageController } from './on-home-page';
 import { recoverAccounts } from './recoverAccounts';
@@ -37,6 +36,11 @@ import type {
   GetAddrFromStarkNameParams,
   GetTransactionStatusParams,
   ListTransactionsParams,
+  AddAccountParams,
+  GetCurrentAccountParams,
+  ListAccountsParams,
+  SwitchAccountParams,
+  ToggleAccountVisibilityParams,
 } from './rpcs';
 import {
   displayPrivateKey,
@@ -53,6 +57,11 @@ import {
   getAddrFromStarkName,
   getTransactionStatus,
   listTransactions,
+  addAccount,
+  getCurrentAccount,
+  listAccounts,
+  switchAccount,
+  toggleAccountVisibility,
 } from './rpcs';
 import { signDeployAccountTransaction } from './signDeployAccountTransaction';
 import type {
@@ -173,7 +182,9 @@ export const onRpcRequest: OnRpcRequestHandler = async ({
         );
 
       case RpcMethod.ListAccounts:
-        return await getStoredUserAccounts(apiParams);
+        return await listAccounts.execute(
+          requestParams as unknown as ListAccountsParams,
+        );
 
       case RpcMethod.DisplayPrivateKey:
         return await displayPrivateKey.execute(
@@ -288,6 +299,26 @@ export const onRpcRequest: OnRpcRequestHandler = async ({
           apiParams.requestParams as unknown as GetAddrFromStarkNameParams,
         );
 
+      case RpcMethod.AddAccount:
+        return await addAccount.execute(
+          requestParams as unknown as AddAccountParams,
+        );
+
+      case RpcMethod.GetCurrentAccount:
+        return await getCurrentAccount.execute(
+          requestParams as unknown as GetCurrentAccountParams,
+        );
+
+      case RpcMethod.SwitchAccount:
+        return await switchAccount.execute(
+          requestParams as unknown as SwitchAccountParams,
+        );
+
+      case RpcMethod.ToggleAccountVisibility:
+        return await toggleAccountVisibility.execute(
+          requestParams as unknown as ToggleAccountVisibilityParams,
+        );
+
       default:
         throw new MethodNotFoundError() as unknown as Error;
     }

packages/starknet-snap/src/on-home-page.test.ts

@@ -1,81 +1,50 @@
 import { ethers } from 'ethers';
-import { constants } from 'starknet';
 
-import { generateAccounts, type StarknetAccount } from './__tests__/helper';
 import { HomePageController } from './on-home-page';
-import type { Network, SnapState } from './types/snapState';
+import { setupAccountController } from './rpcs/__tests__/helper';
+import type { Network } from './types/snapState';
 import {
   BlockIdentifierEnum,
   ETHER_MAINNET,
-  STARKNET_SEPOLIA_TESTNET_NETWORK,
+  STARKNET_MAINNET_NETWORK,
 } from './utils/constants';
 import { loadLocale } from './utils/locale';
-import * as snapHelper from './utils/snap';
 import * as starknetUtils from './utils/starknetUtils';
 
 jest.mock('./utils/snap');
 jest.mock('./utils/logger');
 
 describe('homepageController', () => {
-  const state: SnapState = {
-    accContracts: [],
-    erc20Tokens: [],
-    networks: [STARKNET_SEPOLIA_TESTNET_NETWORK],
-    transactions: [],
-    currentNetwork: STARKNET_SEPOLIA_TESTNET_NETWORK,
-  };
-
-  const mockAccount = async (chainId: constants.StarknetChainId) => {
-    return (await generateAccounts(chainId, 1))[0];
-  };
-
-  const mockState = async () => {
-    const getStateDataSpy = jest.spyOn(snapHelper, 'getStateData');
-    getStateDataSpy.mockResolvedValue(state);
-    return {
-      getStateDataSpy,
-    };
-  };
+  const currentNetwork = STARKNET_MAINNET_NETWORK;
 
   class MockHomePageController extends HomePageController {
-    async getAddress(network: Network): Promise<string> {
-      return super.getAddress(network);
-    }
-
     async getBalance(network: Network, address: string): Promise<string> {
       return super.getBalance(network, address);
     }
   }
 
   describe('execute', () => {
-    const prepareExecuteMock = (account: StarknetAccount, balance: string) => {
-      const getAddressSpy = jest.spyOn(
-        MockHomePageController.prototype,
-        'getAddress',
-      );
+    const setupExecuteTest = async (network: Network, balance = '1000') => {
+      const { account } = await setupAccountController({ network });
+
       const getBalanceSpy = jest.spyOn(
         MockHomePageController.prototype,
         'getBalance',
       );
-      getAddressSpy.mockResolvedValue(account.address);
       getBalanceSpy.mockResolvedValue(balance);
+
       return {
-        getAddressSpy,
+        account,
         getBalanceSpy,
       };
     };
 
     it('returns the correct homepage response', async () => {
       await loadLocale();
-      const { currentNetwork } = state;
-      await mockState();
-      const account = await mockAccount(
-        currentNetwork?.chainId as unknown as constants.StarknetChainId,
-      );
       const balance = '100';
 
-      const { getAddressSpy, getBalanceSpy } = prepareExecuteMock(
-        account,
+      const { getBalanceSpy, account } = await setupExecuteTest(
+        currentNetwork,
         balance,
       );
 
@@ -121,20 +90,15 @@ describe('homepageController', () => {
           type: 'panel',
         },
       });
-      expect(getAddressSpy).toHaveBeenCalledWith(currentNetwork);
       expect(getBalanceSpy).toHaveBeenCalledWith(
         currentNetwork,
         account.address,
       );
     });
 
     it('throws `Failed to initialize Snap HomePage` error if an error was thrown', async () => {
-      await mockState();
-      const account = await mockAccount(constants.StarknetChainId.SN_SEPOLIA);
-      const balance = '100';
-
-      const { getAddressSpy } = prepareExecuteMock(account, balance);
-      getAddressSpy.mockReset().mockRejectedValue(new Error('error'));
+      const { getBalanceSpy } = await setupExecuteTest(currentNetwork);
+      getBalanceSpy.mockReset().mockRejectedValue(new Error('error'));
 
       const homepageController = new MockHomePageController();
       await expect(homepageController.execute()).rejects.toThrow(
@@ -143,84 +107,33 @@ describe('homepageController', () => {
     });
   });
 
-  describe('getAddress', () => {
-    const prepareGetAddressMock = async (account: StarknetAccount) => {
-      const getKeysFromAddressSpy = jest.spyOn(
-        starknetUtils,
-        'getKeysFromAddressIndex',
-      );
-
-      getKeysFromAddressSpy.mockResolvedValue({
-        privateKey: account.privateKey,
-        publicKey: account.publicKey,
-        addressIndex: account.addressIndex,
-        derivationPath: account.derivationPath as unknown as any,
-      });
-
-      const getCorrectContractAddressSpy = jest.spyOn(
-        starknetUtils,
-        'getCorrectContractAddress',
-      );
-      getCorrectContractAddressSpy.mockResolvedValue({
-        address: account.address,
-        signerPubKey: account.publicKey,
-        upgradeRequired: false,
-        deployRequired: false,
-      });
-      return {
-        getKeysFromAddressSpy,
-        getCorrectContractAddressSpy,
-      };
-    };
-
-    it('returns the correct homepage response', async () => {
-      const network = STARKNET_SEPOLIA_TESTNET_NETWORK;
-      await mockState();
-      const account = await mockAccount(constants.StarknetChainId.SN_SEPOLIA);
-      const { getKeysFromAddressSpy, getCorrectContractAddressSpy } =
-        await prepareGetAddressMock(account);
-
-      const homepageController = new MockHomePageController();
-      const result = await homepageController.getAddress(network);
-
-      expect(result).toStrictEqual(account.address);
-      expect(getKeysFromAddressSpy).toHaveBeenCalledWith(
-        // BIP44 Deriver has mocked as undefined, hence this argument should be undefined
-        undefined,
-        network.chainId,
-        state,
-        0,
-      );
-      expect(getCorrectContractAddressSpy).toHaveBeenCalledWith(
-        network,
-        account.publicKey,
-      );
-    });
-  });
-
   describe('getBalance', () => {
-    const prepareGetBalanceMock = async (balance: number) => {
+    const setupGetBalanceTest = async (network: Network, balance: number) => {
+      const { account } = await setupAccountController({ network });
+
       const getBalanceSpy = jest.spyOn(starknetUtils, 'getBalance');
 
       getBalanceSpy.mockResolvedValue(balance.toString(16));
 
       return {
+        account,
         getBalanceSpy,
       };
     };
 
     it('returns the balance on pending block', async () => {
-      const network = STARKNET_SEPOLIA_TESTNET_NETWORK;
       const token = ETHER_MAINNET;
       const expectedBalance = 100;
-      await mockState();
-      const { address } = await mockAccount(
-        constants.StarknetChainId.SN_SEPOLIA,
+      const { getBalanceSpy, account } = await setupGetBalanceTest(
+        currentNetwork,
+        expectedBalance,
       );
-      const { getBalanceSpy } = await prepareGetBalanceMock(expectedBalance);
 
       const homepageController = new MockHomePageController();
-      const result = await homepageController.getBalance(network, address);
+      const result = await homepageController.getBalance(
+        currentNetwork,
+        account.address,
+      );
 
       expect(result).toStrictEqual(
         ethers.utils.formatUnits(
@@ -229,9 +142,9 @@ describe('homepageController', () => {
         ),
       );
       expect(getBalanceSpy).toHaveBeenCalledWith(
-        address,
+        account.address,
         token.address,
-        network,
+        currentNetwork,
         BlockIdentifierEnum.Pending,
       );
     });