test(e2e): speed up manual wallet balance reads (#3030)

tests/src/cucumber/steps/manual_nodes/steps.rs

@@ -979,10 +979,6 @@ async fn step_all_nodes_agree_on_lib(
 
 #[when("I wait for all nodes to be synced to the chain")]
 #[then("I wait for all nodes to be synced to the chain")]
-#[expect(
-    clippy::needless_pass_by_ref_mut,
-    reason = "Cucumber step functions require the world as the first `&mut` argument"
-)]
 async fn step_wait_for_all_nodes_to_be_synced_to_the_chain(
     world: &mut CucumberWorld,
     step: &Step,

tests/src/cucumber/steps/manual_nodes/utils.rs

@@ -38,6 +38,7 @@ use crate::cucumber::{
         display_last_path_components, extract_child_dir_name, funding_wallet_pk_from_node_yaml,
         matching_child_dirs, peer_id_from_node_yaml, track_progress, truncate_hash,
     },
+    wallet::sync::current_wallet_states_for_wallets,
     world::{
         ChainInfoMap, ConfigOverride, CucumberWorld, ManualNodeConfigOverrides, NodeInfo,
         PublicCryptarchiaEndpointPeer, WalletInfo, WalletInfoMap, WalletType,
@@ -308,7 +309,7 @@ pub(crate) fn ensure_fee_sponsorship_and_fork_groups_are_not_mixed(
 }
 
 pub(crate) async fn wait_for_all_nodes_to_be_synced_to_chain(
-    world: &CucumberWorld,
+    world: &mut CucumberWorld,
     step: &str,
 ) -> StepResult {
     let public_cryptarchia_endpoint_peers = world
@@ -347,6 +348,9 @@ pub(crate) async fn wait_for_all_nodes_to_be_synced_to_chain(
                 "All nodes synced to the chain in {:.2?}",
                 start.elapsed()
             );
+
+            catch_up_known_wallet_tracking_after_chain_sync(world, step).await?;
+
             return Ok(());
         }
 
@@ -365,6 +369,27 @@ pub(crate) async fn wait_for_all_nodes_to_be_synced_to_chain(
     }
 }
 
+async fn catch_up_known_wallet_tracking_after_chain_sync(
+    world: &mut CucumberWorld,
+    step: &str,
+) -> StepResult {
+    let wallets = world.wallet_info.values().cloned().collect::<Vec<_>>();
+    if wallets.is_empty() {
+        return Ok(());
+    }
+
+    let started_at = Instant::now();
+    current_wallet_states_for_wallets(world, step, &wallets).await?;
+
+    info!(
+        target: TARGET,
+        "Wallet state refreshed after chain sync in {:.2?}",
+        started_at.elapsed()
+    );
+
+    Ok(())
+}
+
 pub(crate) fn parse_url(raw: &str) -> Result<String, String> {
     let mut trimmed = raw.trim();
     trimmed = trimmed.trim_end_matches('/');

tests/src/cucumber/steps/manual_transactions/command_file_utils.rs

@@ -49,7 +49,7 @@ use tokio::time::{Instant, sleep};
 use tracing::{info, warn};
 
 use crate::{
-    common::wallet::WalletUtxos,
+    common::wallet::{WalletStateView, WalletUtxos},
     cucumber::{
         error::{StepError, StepResult},
         steps::{
@@ -587,8 +587,19 @@ async fn log_wallet_balances(
     step: &str,
     wallets: Vec<WalletInfo>,
 ) -> StepResult {
-    for wallet in wallets {
-        log_wallet_balance(world, step, &wallet.wallet_name).await?;
+    let states = utils::current_wallet_states_for_wallets(world, step, &wallets).await?;
+
+    for wallet in &wallets {
+        let state =
+            states
+                .get(wallet.wallet_name.as_str())
+                .ok_or_else(|| StepError::LogicalError {
+                    message: format!(
+                        "Wallet `{}` balance state is not tracked",
+                        wallet.wallet_name
+                    ),
+                })?;
+        log_wallet_state_balance(&wallet.wallet_name, state);
     }
 
     Ok(())
@@ -599,16 +610,14 @@ async fn log_wallet_balance(
     step: &str,
     wallet_name: &str,
 ) -> StepResult {
-    let available =
-        utils::current_wallet_balance(world, step, wallet_name, WalletOutputState::Available)
-            .await?;
+    let wallet = world.resolve_wallet(wallet_name)?;
+    log_wallet_balances(world, step, vec![wallet]).await
+}
 
-    let reserved =
-        utils::current_wallet_balance(world, step, wallet_name, WalletOutputState::Reserved)
-            .await?;
-
-    let on_chain =
-        utils::current_wallet_balance(world, step, wallet_name, WalletOutputState::OnChain).await?;
+fn log_wallet_state_balance(wallet_name: &str, state: &WalletStateView) {
+    let available = state.balance(WalletOutputState::Available);
+    let reserved = state.balance(WalletOutputState::Reserved);
+    let on_chain = state.balance(WalletOutputState::OnChain);
 
     info!(
         target: TARGET,
@@ -621,8 +630,6 @@ async fn log_wallet_balance(
         on_chain.output_count,
         on_chain.value,
     );
-
-    Ok(())
 }
 
 fn clear_wallet_encumbrances(

tests/src/cucumber/steps/manual_transactions/utils.rs

@@ -17,7 +17,7 @@ pub use crate::cucumber::wallet::{
     sync::{
         current_available_utxos_for_all_wallets, current_available_utxos_for_funding_wallets,
         current_available_utxos_for_user_wallets, current_available_utxos_for_wallet,
-        current_wallet_available_state, current_wallet_balance,
+        current_wallet_available_state, current_wallet_balance, current_wallet_states_for_wallets,
     },
 };
 pub(crate) use crate::cucumber::wallet::{

tests/src/cucumber/wallet/sync.rs

@@ -106,6 +106,17 @@ pub async fn current_wallet_balance(
         .balance(wallet_state_type))
 }
 
+pub async fn current_wallet_states_for_wallets(
+    world: &mut CucumberWorld,
+    step: &str,
+    wallets: &[WalletInfo],
+) -> Result<BTreeMap<WalletId, WalletStateView>, StepError> {
+    let feed_requirements = wallet_feed_source_requirements(world, wallets).await?;
+    let wallet_keys = build_tracked_wallet_keys(world, step, wallets)?;
+
+    current_wallet_state_views(world, &wallet_keys, feed_requirements).await
+}
+
 pub async fn current_wallet_available_state(
     world: &mut CucumberWorld,
     wallet_name: &str,