fix: unshiedling

apps/extension/src/Setup/AccountCreation/SeedPhrase.tsx

@@ -57,6 +57,7 @@ export const SeedPhrase: React.FC<Props> = (props) => {
             words={seedPhrase}
           />
           <FeedbackButton
+            data-testid="setup-copy-to-clipboard-button"
             className="text-center mx-auto block"
             successMessage="Copied to clipboard"
             errorMessage="Error trying to copy"

apps/extension/src/Setup/AccountCreation/SeedPhraseConfirmation.tsx

@@ -100,7 +100,7 @@ export const SeedPhraseConfirmation = (
           className="[&_label]:!text-sm justify-center flex-1"
         >
           <Input
-            data-testd="setup-seed-phrase-verification-1-input"
+            data-testid="setup-seed-phrase-verification-1-input"
             label={`Word #${index1ToConfirm + 1}`}
             value={verificationInput1 || ""}
             error={index1Error}

apps/extension/src/Setup/AccountCreation/SeedPhraseWarning.tsx

@@ -41,7 +41,7 @@ export const SeedPhraseWarning = ({
       </Stack>
       <footer>
         <ActionButton
-          data-testid="setup-show-phrase-button"
+          data-testid={countdown === 0 && "setup-show-phrase-button"}
           size="lg"
           className={clsx({ "pointer-events-none opacity-50": countdown > 0 })}
           onClick={onComplete}

apps/extension/src/Setup/Common/CreateKeyForm.tsx

@@ -42,7 +42,7 @@ export default function CreateKeyForm({
     >
       <Stack gap={12} className="justify-center flex-1">
         <AccountAlias
-          data-testid="setup-seedphrase-alias-input"
+          data-testid="setup-seed-phrase-alias-input"
           value={accountName}
           onChange={setAccountName}
         />

apps/extension/src/background/keyring/handler.ts

@@ -5,6 +5,7 @@ import {
   QueryAccountsMsg,
   QueryBalancesMsg,
   QueryDefaultAccountMsg,
+  ShieldedSyncMsg,
   VerifyArbitraryMsg,
 } from "provider/messages";
 import { Env, Handler, InternalHandler, Message } from "router";
@@ -60,6 +61,8 @@ export const getHandler: (service: KeyRingService) => Handler = (service) => {
         return handleQueryAccountsMsg(service)(env, msg as QueryAccountsMsg);
       case QueryBalancesMsg:
         return handleQueryBalancesMsg(service)(env, msg as QueryBalancesMsg);
+      case ShieldedSyncMsg:
+        return handleShieldedSyncMsg(service)(env, msg as ShieldedSyncMsg);
       case SetActiveAccountMsg:
         return handleSetActiveAccountMsg(service)(
           env,
@@ -230,6 +233,14 @@ const handleQueryBalancesMsg: (
   };
 };
 
+const handleShieldedSyncMsg: (
+  service: KeyRingService
+) => InternalHandler<ShieldedSyncMsg> = (service) => {
+  return async (_, { addresses }) => {
+    return await service.shieldedSync(addresses);
+  };
+};
+
 const handleQueryParentAccountsMsg: (
   service: KeyRingService
 ) => InternalHandler<QueryParentAccountsMsg> = (service) => {

apps/extension/src/background/keyring/init.ts

@@ -5,6 +5,7 @@ import {
   QueryAccountsMsg,
   QueryBalancesMsg,
   QueryDefaultAccountMsg,
+  ShieldedSyncMsg,
   VerifyArbitraryMsg,
 } from "provider/messages";
 import { Router } from "router";
@@ -38,6 +39,7 @@ export function init(router: Router, service: KeyRingService): void {
   router.registerMessage(QueryAccountsMsg);
   router.registerMessage(QueryDefaultAccountMsg);
   router.registerMessage(QueryBalancesMsg);
+  router.registerMessage(ShieldedSyncMsg);
   router.registerMessage(QueryParentAccountsMsg);
   router.registerMessage(SaveAccountSecretMsg);
   router.registerMessage(ScanAccountsMsg);

apps/extension/src/background/keyring/service.ts

@@ -452,6 +452,21 @@ export class KeyRingService {
     return Sdk.has_masp_params();
   }
 
+  async shieldedSync(addresses: string[]): Promise<void> {
+    const vks = [];
+    for await (const address of addresses) {
+      const account = await this.vaultService.findOneOrFail<AccountStore>(
+        KEYSTORE_KEY,
+        "address",
+        address
+      );
+      vks.push(account.public.owner);
+    }
+
+    const query = await this.sdkService.getQuery();
+    await query.shielded_sync(vks);
+  }
+
   async queryBalances(
     owner: string,
     tokens: string[]

apps/extension/src/provider/InjectedNamada.ts

@@ -60,6 +60,13 @@ export class InjectedNamada implements INamada {
     >("balances", props);
   }
 
+  public async shieldedSync(props: { addresses: string[] }): Promise<void> {
+    return await InjectedProxy.requestMethod<{ addresses: string[] }, void>(
+      "shieldedSync",
+      props
+    );
+  }
+
   public async getChain(): Promise<Chain | undefined> {
     return await InjectedProxy.requestMethod<void, Chain>("getChain");
   }

apps/extension/src/provider/Namada.ts

@@ -22,6 +22,7 @@ import {
   QueryAccountsMsg,
   QueryBalancesMsg,
   QueryDefaultAccountMsg,
+  ShieldedSyncMsg,
   VerifyArbitraryMsg,
 } from "./messages";
 
@@ -125,6 +126,14 @@ export class Namada implements INamada {
     );
   }
 
+  public async shieldedSync(props: { addresses: string[] }): Promise<void> {
+    const { addresses } = props;
+    return await this.requester?.sendMessage(
+      Ports.Background,
+      new ShieldedSyncMsg(addresses)
+    );
+  }
+
   public async submitTx(props: TxMsgProps): Promise<void> {
     return await this.requester?.sendMessage(
       Ports.Background,

apps/extension/src/provider/messages.ts

@@ -27,6 +27,7 @@ enum MessageType {
   QueryDefaultAccount = "query-default-account",
   ApproveTx = "approve-tx",
   QueryBalances = "query-balances",
+  ShieldedSync = "shielded-sync",
   SubmitIbcTransfer = "submit-ibc-transfer",
   SubmitLedgerTransfer = "submit-ledger-transfer",
   EncodeRevealPublicKey = "encode-reveal-public-key",
@@ -167,6 +168,28 @@ export class QueryBalancesMsg extends Message<
   }
 }
 
+export class ShieldedSyncMsg extends Message<void> {
+  public static type(): MessageType {
+    return MessageType.ShieldedSync;
+  }
+
+  constructor(public readonly addresses: string[]) {
+    super();
+  }
+
+  validate(): void {
+    validateProps(this, ["addresses"]);
+  }
+
+  route(): string {
+    return Route.KeyRing;
+  }
+
+  type(): string {
+    return ShieldedSyncMsg.type();
+  }
+}
+
 export class QueryDefaultAccountMsg extends Message<
   DerivedAccount | undefined
 > {

apps/namada-interface/src/slices/accounts.ts

@@ -3,7 +3,7 @@ import BigNumber from "bignumber.js";
 import { atom } from "jotai";
 
 import { chains } from "@namada/chains";
-import { getIntegration } from "@namada/integrations";
+import { getIntegration, Integration } from "@namada/integrations";
 import { Account as AccountDetails, ChainKey, TokenType } from "@namada/types";
 
 import { chainAtom } from "slices/chain";
@@ -193,18 +193,55 @@ const balancesAtom = (() => {
 
       set(base, {});
 
-      accounts.forEach(async (account) => {
-        const result = await namada.queryBalances(account.address, [
-          nativeToken || tokenAddress,
-        ]);
-        const balance = result.reduce(
-          (acc, curr) => ({ ...acc, [curr.token]: new BigNumber(curr.amount) }),
-          {} as Balance
-        );
-        set(base, { ...get(base), [account.address]: balance });
+      // Split accounts into transparent and shielded
+      const [transparentAccounts, shieldedAccounts] = accounts.reduce((acc, curr) => {
+        if (curr.isShielded) {
+          acc[1].push(curr)
+        } else {
+          acc[0].push(curr)
+        }
+        return acc;
+      }, [[], []] as [AccountDetails[], AccountDetails[]]);
+
+      const token = nativeToken || tokenAddress;
+
+      // We query the balances for the transparent accounts first as it's faster
+      const transparentBalances = await Promise.all(queryBalance(namada, transparentAccounts, token));
+      transparentBalances.forEach(([address, balance ]) => {
+        set(base, { ...get(base), [address]: balance });
       });
+
+
+      // Before querying the shielded balances we need to sync the shielded context
+      const shieldedAddresses = accounts
+        .filter((account) => account.isShielded)
+        .map((account) => account.address);
+      await namada.sync(shieldedAddresses);
+
+      const shieldedBalances = await Promise.all(queryBalance(namada, shieldedAccounts, token));
+      shieldedBalances.forEach(([address, balance ]) => {
+        set(base, { ...get(base), [address]: balance });
+      });
+
     }
   );
 })();
 
+const queryBalance = (
+  int: InstanceType<Integration>,
+  accounts: AccountDetails[],
+  token: string,
+): Promise<readonly [string, Balance]>[] => {
+  return accounts.map(async (account) => {
+    const result = await int.queryBalances(account.address, [token]);
+    return [
+      account.address,
+      result.reduce(
+        (acc, curr) => ({ ...acc, [curr.token]: new BigNumber(curr.amount) }),
+        {} as Balance,
+      ),
+    ] as const;
+  });
+};
+
 export { accountsAtom, balancesAtom };

e2e/setup-namada.sh

@@ -1,6 +1,6 @@
 #!/bin/bash -x
 
-VERSION="v0.31.6"
+VERSION="v0.31.8"
 CURRENT_VERSION=""
 NAMADA_DIR=".namada"
 NAMADA_BASE_DIR=".namada/basedir"

e2e/src/partial/setup.ts

@@ -50,20 +50,24 @@ export const importAccount = async (
   await page.waitForNavigation();
 
   // Fill alias and pwd
-  await (
-    await page.$("[data-testid='setup-import-keys-alias-input']")
-  )?.type(alias0);
+  const aliasInput = await page.$(
+    "[data-testid='setup-seed-phrase-alias-input'] input"
+  );
+  await aliasInput?.type(alias0);
 
   const pwdInputs = await page.$$(
-    "[data-testid='setup-import-keys-pwd-input']"
+    "[data-testid='setup-seed-phrase-pwd-input'] input"
   );
+
   for await (const input of pwdInputs) {
     await input.type(pwd);
   }
 
   // Click on next button
   (
-    await page.waitForSelector("[data-testid='setup-import-keys-next-button']")
+    await page.waitForSelector(
+      "[data-testid='setup-seed-phrase-verification-next-btn']"
+    )
   )?.click();
   await page.waitForNavigation();
 
@@ -92,7 +96,9 @@ export const createAccount = async (
 
   // Click on show phrase
   (
-    await page.waitForSelector("[data-testid='setup-show-phrase-button']")
+    await page.waitForSelector("[data-testid='setup-show-phrase-button']", {
+      timeout: 5500,
+    })
   )?.click();
 
   await page.waitForNavigation();
@@ -107,7 +113,7 @@ export const createAccount = async (
 
   const words = await page.$$eval(
     "[data-testid='setup-seed-phrase-list'] li",
-    (els) => els.map((el) => el.textContent as string)
+    (els) => els.map((el) => (el.textContent as string).replace(/\d+\s/, ""))
   );
 
   expect(await page.evaluate(() => navigator.clipboard.readText())).toEqual(
@@ -123,11 +129,11 @@ export const createAccount = async (
   await page.waitForNavigation();
 
   // Verify mnemonic
-  const input1 = await page.$(
+  const input1 = await page.waitForSelector(
     "[data-testid='setup-seed-phrase-verification-1-input']"
   );
 
-  const input2 = await page.$(
+  const input2 = await page.waitForSelector(
     "[data-testid='setup-seed-phrase-verification-2-input']"
   );
 
@@ -144,14 +150,21 @@ export const createAccount = async (
   await (await input1?.$("input"))?.type(words[Number(word1Number) - 1]);
   await (await input2?.$("input"))?.type(words[Number(word2Number) - 1]);
 
+  (
+    await page.waitForSelector(
+      "[data-testid='setup-seed-phrase-verification-next-btn']"
+    )
+  )?.click();
+  await page.waitForNavigation();
+
   // Fill alias and pwd
   const aliasInput = await page.$(
-    "[data-testid='setup-seed-phrase-alias-input']"
+    "[data-testid='setup-seed-phrase-alias-input'] input"
   );
   await aliasInput?.type(alias0);
 
   const pwdInputs = await page.$$(
-    "[data-testid='setup-seed-phrase-pwd-input']"
+    "[data-testid='setup-seed-phrase-pwd-input'] input"
   );
 
   for await (const input of pwdInputs) {
@@ -163,7 +176,6 @@ export const createAccount = async (
       "[data-testid='setup-seed-phrase-verification-next-btn']"
     )
   )?.click();
-
   await page.waitForNavigation();
 
   const closePageBtn = await page.waitForSelector(

packages/components/src/FeedbackButton.tsx

@@ -8,6 +8,7 @@ type Props = {
   children: React.ReactNode;
   successMessage: string;
   errorMessage: string;
+  "data-testid"?: string;
   onAction: () => void;
 } & React.ComponentPropsWithoutRef<"button">;
 
@@ -16,6 +17,7 @@ export const FeedbackButton = ({
   successMessage,
   errorMessage,
   onAction,
+  "data-testid": dataTestid,
   ...props
 }: Props): JSX.Element => {
   const [displaySuccess, setDisplaySuccess] = useState(false);
@@ -49,6 +51,7 @@ export const FeedbackButton = ({
   return (
     <div className="relative">
       <button
+        data-testid={dataTestid}
         onClick={onPerformClick}
         className={clsx(
           "text-white font-medium underline text-xs hover:text-yellow",

packages/integrations/src/Keplr.ts

@@ -46,7 +46,7 @@ class Keplr implements Integration<Account, OfflineSigner> {
    * override keplr instance for testing
    * @param chain
    */
-  constructor(public readonly chain: Chain) { }
+  constructor(public readonly chain: Chain) {}
 
   private init(): void {
     if (!this._keplr) {
@@ -221,6 +221,9 @@ class Keplr implements Integration<Account, OfflineSigner> {
         };
       });
   }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  public async sync(_addresses: string[]): Promise<void> {}
 }
 
 export default Keplr;