Merge branch 'tomas/fix-balance-query' (#1751)

* origin/tomas/fix-balance-query:
  changelog: add #1751
  apps/client/rpc: fix transparent balance query without explicit token
  test/e2e/tx_and_queries: add check for balance without explicit token
  apps/client/utils/init_network: remove redundant add address call
  apps/wallet: fix add genesis addresses to add token VP types
  apps/client/rpc: replace lookup alias with method from wallet
  shared/wallet: add fns to lookup aliases and tokens
  apps/cli/context: rm unused functions

.changelog/unreleased/bug-fixes/1751-fix-balance-query.md

@@ -0,0 +1,2 @@
+- Fixed transparent balance query when only an owner address is specified without
+  an explicit token. ([\#1751](https://github.com/anoma/namada/pull/1751))

apps/src/lib/cli/context.rs

@@ -1,14 +1,12 @@
 //! CLI input types can be used for command arguments
 
-use std::collections::{HashMap, HashSet};
 use std::env;
 use std::marker::PhantomData;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
 use color_eyre::eyre::Result;
 use namada::ledger::masp::ShieldedContext;
-use namada::ledger::wallet::store::AddressVpType;
 use namada::ledger::wallet::Wallet;
 use namada::types::address::Address;
 use namada::types::chain::ChainId;
@@ -193,32 +191,6 @@ impl Context {
     pub fn read_wasm(&self, file_name: impl AsRef<Path>) -> Vec<u8> {
         wasm_loader::read_wasm_or_exit(self.wasm_dir(), file_name)
     }
-
-    /// Try to find an alias for a given address from the wallet. If not found,
-    /// formats the address into a string.
-    pub fn lookup_alias(&self, addr: &Address) -> String {
-        match self.wallet.find_alias(addr) {
-            Some(alias) => alias.to_string(),
-            None => addr.to_string(),
-        }
-    }
-
-    /// Get addresses with tokens VP type.
-    pub fn tokens(&self) -> HashSet<Address> {
-        self.wallet.get_addresses_with_vp_type(AddressVpType::Token)
-    }
-
-    /// Get addresses with tokens VP type associated with their aliases.
-    pub fn tokens_with_aliases(&self) -> HashMap<Address, String> {
-        self.wallet
-            .get_addresses_with_vp_type(AddressVpType::Token)
-            .into_iter()
-            .map(|addr| {
-                let alias = self.lookup_alias(&addr);
-                (addr, alias)
-            })
-            .collect()
-    }
 }
 
 /// Load global config from expected path in the `base_dir` or try to generate a

apps/src/lib/client/rpc.rs

@@ -193,7 +193,7 @@ pub async fn query_transfers<
         for (account, MaspChange { ref asset, change }) in tfer_delta {
             if account != masp() {
                 print!("  {}:", account);
-                let token_alias = lookup_alias(wallet, asset);
+                let token_alias = wallet.lookup_alias(asset);
                 let sign = match change.cmp(&Change::zero()) {
                     Ordering::Greater => "+",
                     Ordering::Less => "-",
@@ -216,7 +216,7 @@ pub async fn query_transfers<
             if fvk_map.contains_key(&account) {
                 print!("  {}:", fvk_map[&account]);
                 for (token_addr, val) in masp_change {
-                    let token_alias = lookup_alias(wallet, &token_addr);
+                    let token_alias = wallet.lookup_alias(&token_addr);
                     let sign = match val.cmp(&Change::zero()) {
                         Ordering::Greater => "+",
                         Ordering::Less => "-",
@@ -303,11 +303,12 @@ pub async fn query_transparent_balance<
         Address::Internal(namada::types::address::InternalAddress::Multitoken)
             .to_db_key(),
     );
+    let tokens = wallet.tokens_with_aliases();
     match (args.token, args.owner) {
         (Some(token), Some(owner)) => {
             let balance_key =
                 token::balance_key(&token, &owner.address().unwrap());
-            let token_alias = lookup_alias(wallet, &token);
+            let token_alias = wallet.lookup_alias(&token);
             match query_storage_value::<C, token::Amount>(client, &balance_key)
                 .await
             {
@@ -323,17 +324,15 @@ pub async fn query_transparent_balance<
             }
         }
         (None, Some(owner)) => {
-            let balances =
-                query_storage_prefix::<C, token::Amount>(client, &prefix).await;
-            if let Some(balances) = balances {
-                print_balances(
-                    client,
-                    wallet,
-                    balances,
-                    None,
-                    owner.address().as_ref(),
-                )
-                .await;
+            let owner = owner.address().unwrap();
+            for (token_alias, token) in tokens {
+                let balance = get_token_balance(client, &token, &owner).await;
+                if !balance.is_zero() {
+                    let balance =
+                        format_denominated_amount(client, &token, balance)
+                            .await;
+                    println!("{}: {}", token_alias, balance);
+                }
             }
         }
         (Some(token), None) => {
@@ -425,7 +424,7 @@ pub async fn query_pinned_balance<
                 println!("Payment address {} has not yet been consumed.", owner)
             }
             (Ok((balance, epoch)), Some(token)) => {
-                let token_alias = lookup_alias(wallet, token);
+                let token_alias = wallet.lookup_alias(token);
 
                 let total_balance = balance
                     .get(&(epoch, token.clone()))
@@ -513,7 +512,7 @@ async fn print_balances<C: namada::ledger::queries::Client + Sync>(
                 format!(
                     ": {}, owned by {}",
                     format_denominated_amount(client, tok, balance).await,
-                    lookup_alias(wallet, owner)
+                    wallet.lookup_alias(owner)
                 ),
             ),
             None => continue,
@@ -540,7 +539,7 @@ async fn print_balances<C: namada::ledger::queries::Client + Sync>(
                 // the token has been already printed
             }
             _ => {
-                let token_alias = lookup_alias(wallet, &t);
+                let token_alias = wallet.lookup_alias(&t);
                 writeln!(w, "Token {}", token_alias).unwrap();
                 print_token = Some(t);
             }
@@ -555,11 +554,11 @@ async fn print_balances<C: namada::ledger::queries::Client + Sync>(
             (Some(_), Some(target)) | (None, Some(target)) => writeln!(
                 w,
                 "No balances owned by {}",
-                lookup_alias(wallet, target)
+                wallet.lookup_alias(target)
             )
             .unwrap(),
             (Some(token), None) => {
-                let token_alias = lookup_alias(wallet, token);
+                let token_alias = wallet.lookup_alias(token);
                 writeln!(w, "No balances for token {}", token_alias).unwrap()
             }
             (None, None) => writeln!(w, "No balances").unwrap(),
@@ -750,7 +749,7 @@ pub async fn query_shielded_balance<
                     .expect("context should contain viewing key")
             };
 
-            let token_alias = lookup_alias(wallet, &token);
+            let token_alias = wallet.lookup_alias(&token);
 
             let total_balance = balance
                 .get(&(epoch, token.clone()))
@@ -847,10 +846,10 @@ pub async fn query_shielded_balance<
                     .as_ref(),
             )
             .unwrap();
-            let token_alias = lookup_alias(wallet, &token);
+            let token_alias = wallet.lookup_alias(&token);
             println!("Shielded Token {}:", token_alias);
             let mut found_any = false;
-            let token_alias = lookup_alias(wallet, &token);
+            let token_alias = wallet.lookup_alias(&token);
             println!("Shielded Token {}:", token_alias,);
             for fvk in viewing_keys {
                 // Query the multi-asset balance at the given spending key
@@ -928,7 +927,7 @@ pub async fn print_decoded_balance<
         {
             println!(
                 "{} : {}",
-                lookup_alias(wallet, token_addr),
+                wallet.lookup_alias(token_addr),
                 format_denominated_amount(client, token_addr, (*amount).into())
                     .await,
             );
@@ -2290,15 +2289,6 @@ pub async fn get_governance_parameters<
     namada::ledger::rpc::get_governance_parameters(client).await
 }
 
-/// Try to find an alias for a given address from the wallet. If not found,
-/// formats the address into a string.
-fn lookup_alias(wallet: &Wallet<CliWalletUtils>, addr: &Address) -> String {
-    match wallet.find_alias(addr) {
-        Some(alias) => format!("{}", alias),
-        None => format!("{}", addr),
-    }
-}
-
 /// A helper to unwrap client's response. Will shut down process on error.
 fn unwrap_client_response<C: namada::ledger::queries::Client, T>(
     response: Result<T, C::Error>,

apps/src/lib/client/utils.rs

@@ -645,11 +645,10 @@ pub fn init_network(
         })
     }
 
-    config.token.iter_mut().for_each(|(name, config)| {
+    config.token.iter_mut().for_each(|(_name, config)| {
         if config.address.is_none() {
             let address = address::gen_established_address("token");
             config.address = Some(address.to_string());
-            wallet.add_address(name.clone(), address, true);
         }
         if config.vp.is_none() {
             config.vp = Some("vp_token".to_string());

apps/src/lib/wallet/mod.rs

@@ -5,14 +5,17 @@ mod store;
 
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
+use std::str::FromStr;
 use std::{env, fs};
 
 use namada::bip39::{Language, Mnemonic};
 pub use namada::ledger::wallet::alias::Alias;
 use namada::ledger::wallet::{
-    ConfirmationResponse, FindKeyError, GenRestoreKeyError, Wallet, WalletUtils,
+    AddressVpType, ConfirmationResponse, FindKeyError, GenRestoreKeyError,
+    Wallet, WalletUtils,
 };
 pub use namada::ledger::wallet::{ValidatorData, ValidatorKeys};
+use namada::types::address::Address;
 use namada::types::key::*;
 use rand_core::OsRng;
 pub use store::wallet_file;
@@ -237,6 +240,13 @@ pub fn add_genesis_addresses(
     wallet: &mut Wallet<CliWalletUtils>,
     genesis: GenesisConfig,
 ) {
+    for (alias, config) in &genesis.token {
+        let exp_addr = format!("Genesis token {alias} must have address");
+        let address =
+            Address::from_str(config.address.as_ref().expect(&exp_addr))
+                .expect("Valid address");
+        wallet.add_vp_type_to_address(AddressVpType::Token, address);
+    }
     for (alias, addr) in defaults::addresses_from_genesis(genesis) {
         wallet.add_address(alias.normalize(), addr, true);
     }

shared/src/ledger/wallet/mod.rs

@@ -5,7 +5,7 @@ mod keys;
 pub mod pre_genesis;
 pub mod store;
 
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::fmt::Display;
 use std::str::FromStr;
 
@@ -503,6 +503,15 @@ impl<U: WalletUtils> Wallet<U> {
         self.store.find_alias(address)
     }
 
+    /// Try to find an alias for a given address from the wallet. If not found,
+    /// formats the address into a string.
+    pub fn lookup_alias(&self, addr: &Address) -> String {
+        match self.find_alias(addr) {
+            Some(alias) => format!("{}", alias),
+            None => format!("{}", addr),
+        }
+    }
+
     /// Get all known addresses by their alias, paired with PKH, if known.
     pub fn get_addresses(&self) -> HashMap<String, Address> {
         self.store
@@ -678,4 +687,15 @@ impl<U: WalletUtils> Wallet<U> {
     pub fn store_dir(&self) -> &U::Storage {
         &self.store_dir
     }
+
+    /// Get addresses with tokens VP type keyed and ordered by their aliases.
+    pub fn tokens_with_aliases(&self) -> BTreeMap<String, Address> {
+        self.get_addresses_with_vp_type(AddressVpType::Token)
+            .into_iter()
+            .map(|addr| {
+                let alias = self.lookup_alias(&addr);
+                (alias, addr)
+            })
+            .collect()
+    }
 }

tests/src/e2e/ledger_tests.rs

@@ -611,13 +611,41 @@ fn ledger_txs_and_queries() -> Result<()> {
                 &validator_one_rpc,
             ],
             // expect a decimal
-            r"nam: \d+(\.\d+)?",
+            vec![r"nam: \d+(\.\d+)?"],
+        ),
+        // Unspecified token expect all tokens from wallet derived from genesis
+        (
+            vec!["balance", "--owner", BERTHA, "--node", &validator_one_rpc],
+            // expect all genesis tokens, sorted by alias
+            vec![
+                r"apfel: \d+(\.\d+)?",
+                r"btc: \d+(\.\d+)?",
+                r"dot: \d+(\.\d+)?",
+                r"eth: \d+(\.\d+)?",
+                r"kartoffel: \d+(\.\d+)?",
+                r"schnitzel: \d+(\.\d+)?",
+            ],
         ),
     ];
     for (query_args, expected) in &query_args_and_expected_response {
+        // Run as a non-validator
         let mut client = run!(test, Bin::Client, query_args, Some(40))?;
-        client.exp_regex(expected)?;
+        for pattern in expected {
+            client.exp_regex(pattern)?;
+        }
+        client.assert_success();
 
+        // Run as a validator
+        let mut client = run_as!(
+            test,
+            Who::Validator(0),
+            Bin::Client,
+            query_args,
+            Some(40)
+        )?;
+        for pattern in expected {
+            client.exp_regex(pattern)?;
+        }
         client.assert_success();
     }
     let christel = find_address(&test, CHRISTEL)?;