fix: [lw-12156] fix remove wallet account functionality in nami mode (#…

…1672)
input-output-hk · Jan 29, 2025 · 19836a4 · 19836a4
1 parent 2fde1f7
commit 19836a4
Show file tree

Hide file tree

Showing 5 changed files with 218 additions and 109 deletions.
diff --git a/apps/browser-extension-wallet/src/hooks/useWalletManager.ts b/apps/browser-extension-wallet/src/hooks/useWalletManager.ts
@@ -662,7 +662,10 @@ export const useWalletManager = (): UseWalletManager => {
    */
   const deleteWallet = useCallback(
     // eslint-disable-next-line max-statements
-    async (isForgotPasswordFlow = false): Promise<WalletManagerActivateProps | undefined> => {
+    async (
+      isForgotPasswordFlow = false,
+      nextWallet?: AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>
+    ): Promise<WalletManagerActivateProps | undefined> => {
       let walletToDelete: Pick<WalletManagerActivateProps, 'walletId'> = await firstValueFrom(
         walletManager.activeWalletId$
       );
@@ -685,7 +688,7 @@ export const useWalletManager = (): UseWalletManager => {
       }
       await walletRepository.removeWallet(walletToDelete.walletId);
 
-      const wallets = await firstValueFrom(walletRepository.wallets$);
+      const wallets = nextWallet ? [nextWallet] : await firstValueFrom(walletRepository.wallets$);
       if (wallets.length > 0) {
         const activateProps = {
           walletId: wallets[0].walletId,

diff --git a/packages/nami/src/adapters/account.test.ts b/packages/nami/src/adapters/account.test.ts
@@ -36,7 +36,7 @@ const genWallet = (
     walletId,
     type,
     accounts: [acc0, acc1, acc2],
-  }) as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>;
+  }) as unknown as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>;
 
 const wallet1 = genWallet('wallet1', 'InMemory');
 const wallet2 = genWallet('wallet2', 'InMemory');
@@ -318,7 +318,10 @@ describe('useAccount', () => {
               },
             },
           ],
-        } as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>,
+        } as unknown as AnyWallet<
+          Wallet.WalletMetadata,
+          Wallet.AccountMetadata
+        >,
       ]);
     });
 
@@ -364,7 +367,7 @@ describe('useAccount', () => {
       {
         walletId: 'wallet',
         accounts: [acc1],
-      } as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>,
+      } as unknown as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>,
     ]) as Wallets$;
     const activeWalletId$ = of({
       walletId: 'wallet',
@@ -383,12 +386,82 @@ describe('useAccount', () => {
       }),
     );
 
-    await result.current.removeAccount({
+    const accountToRemove = {
       accountIndex: 1,
       walletId: 'wallet',
+    };
+
+    await result.current.removeAccount(accountToRemove);
+
+    expect(mockRemoveWallet).not.toHaveBeenCalled();
+
+    expect(mockRemoveAccount).toHaveBeenCalledWith(accountToRemove);
+  });
+
+  it('should call removeWallet with correct arguments (no next wallet)', async () => {
+    const activeWallet = {
+      walletId: 'wallet',
+      accounts: [acc1],
+    } as unknown as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>;
+    const wallets$ = of([activeWallet]) as Wallets$;
+    const activeWalletId$ = of({
+      walletId: 'wallet',
+      accountIndex: 1,
+    }) as unknown as WalletManagerApi['activeWalletId$'];
+
+    const { result } = renderHook(() =>
+      useAccountUtil({
+        wallets$,
+        activeWalletId$,
+        updateAccountMetadata: mockUpdateAccountMetadata,
+        activateAccount: mockActivateAccount,
+        addAccount: mockAddAccount,
+        removeAccount: mockRemoveAccount,
+        removeWallet: mockRemoveWallet,
+      }),
+    );
+
+    await result.current.removeWallet({
+      nextWalletId: 'wallet1',
+    });
+
+    expect(mockRemoveWallet).toHaveBeenCalledWith(false, undefined);
+
+    expect(mockRemoveAccount).not.toHaveBeenCalled();
+  });
+
+  it('should call removeWallet with correct arguments', async () => {
+    const activeWallet = {
+      walletId: 'wallet',
+      accounts: [acc1],
+    } as unknown as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>;
+    const nextWallet = {
+      walletId: 'wallet1',
+      accounts: [acc0],
+    } as unknown as AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>;
+    const wallets$ = of([activeWallet, nextWallet]) as Wallets$;
+    const activeWalletId$ = of({
+      walletId: 'wallet',
+      accountIndex: 1,
+    }) as unknown as WalletManagerApi['activeWalletId$'];
+
+    const { result } = renderHook(() =>
+      useAccountUtil({
+        wallets$,
+        activeWalletId$,
+        updateAccountMetadata: mockUpdateAccountMetadata,
+        activateAccount: mockActivateAccount,
+        addAccount: mockAddAccount,
+        removeAccount: mockRemoveAccount,
+        removeWallet: mockRemoveWallet,
+      }),
+    );
+
+    await result.current.removeWallet({
+      nextWalletId: 'wallet1',
     });
 
-    expect(mockRemoveWallet).toHaveBeenCalledWith();
+    expect(mockRemoveWallet).toHaveBeenCalledWith(false, nextWallet);
 
     expect(mockRemoveAccount).not.toHaveBeenCalled();
   });

diff --git a/packages/nami/src/adapters/account.ts b/packages/nami/src/adapters/account.ts
@@ -35,7 +35,12 @@ interface AccountsProps {
   removeAccount: (
     props: Readonly<RemoveAccountProps>,
   ) => Promise<RemoveAccountProps>;
-  removeWallet: (isChangePasswordFlow?: boolean) => Promise<void>;
+  removeWallet: (
+    isChangePasswordFlow?: boolean,
+    nextWallet?: Readonly<
+      AnyWallet<Wallet.WalletMetadata, Wallet.AccountMetadata>
+    >,
+  ) => Promise<void>;
   updateAccountMetadata: (
     props: Readonly<UpdateAccountMetadataProps<Wallet.AccountMetadata>>,
   ) => Promise<UpdateAccountMetadataProps<Wallet.AccountMetadata>>;
@@ -49,7 +54,7 @@ export interface Account {
   balance?: Partial<Record<Wallet.ChainName, string>>;
   address?: Partial<Record<Wallet.ChainName, string>>;
   recentSendToAddress?: Partial<Record<Wallet.ChainName, string>>;
-  type?: WalletType;
+  type: WalletType;
 }
 
 export interface UseAccount {
@@ -77,6 +82,11 @@ export interface UseAccount {
       walletId?: WalletId;
     }>,
   ) => Promise<void>;
+  removeWallet: (
+    props: Readonly<{
+      nextWalletId?: WalletId;
+    }>,
+  ) => Promise<void>;
   updateAccountMetadata: (
     data: Readonly<{
       name?: string;
@@ -260,20 +270,26 @@ export const useAccountUtil = ({
       },
       [wallets, addAccount],
     ),
+    removeWallet: useCallback(
+      async ({ nextWalletId }) => {
+        if (walletId === undefined) {
+          return;
+        }
+
+        const wallet = wallets?.find(elm => elm.walletId === nextWalletId);
+        await removeWallet(false, wallet);
+      },
+      [removeWallet, walletId, wallets],
+    ),
     removeAccount: useCallback(
       async ({ accountIndex, walletId }) => {
         if (walletId === undefined) {
           return;
         }
-        const isLastAccount = !allAccountsSorted.some(
-          a => a.walletId === walletId && a.index !== accountIndex,
-        );
 
-        await (isLastAccount
-          ? removeWallet()
-          : removeAccount({ accountIndex, walletId }));
+        await removeAccount({ accountIndex, walletId });
       },
-      [removeAccount, walletId, allAccountsSorted],
+      [removeAccount, walletId],
     ),
     activateAccount: useCallback(
       async props => {