fix: deflake //rs/tests/ckbtc:ckbtc_minter_update_balance_head_nns

rs/tests/ckbtc/ckbtc_minter_update_balance.rs

@@ -198,7 +198,7 @@ pub fn test_update_balance(env: TestEnv) {
             &logger,
             "Calling update balance on second subaccount with missing ledger."
         );
-        assert_temporarily_unavailable(&minter_agent, &subaccount2).await;
+        assert_temporarily_unavailable(&minter_agent, &logger, &subaccount2).await;
 
         // The ledger canister is back online.
         start_canister(&ledger_canister).await;

rs/tests/ckbtc/src/utils.rs

@@ -23,11 +23,12 @@ use crate::{ADDRESS_LENGTH, IcRpcClientType};
 use assert_matches::assert_matches;
 use candid::{Decode, Encode, Nat};
 use canister_test::Canister;
+use ic_agent::AgentError;
 use ic_btc_adapter_test_utils::{
     bitcoin::{Address, Amount, Txid},
     rpc_client::{Auth, RpcClient, RpcClientType},
 };
-use ic_ckbtc_agent::CkBtcMinterAgent;
+use ic_ckbtc_agent::{CkBtcMinterAgent, CkBtcMinterAgentError};
 use ic_ckbtc_minter::state::RetrieveBtcStatus;
 use ic_ckbtc_minter::updates::retrieve_btc::{RetrieveBtcArgs, RetrieveBtcError};
 use ic_ckbtc_minter::updates::update_balance::UtxoStatus::Checked;
@@ -630,25 +631,48 @@ pub async fn assert_no_new_utxo(
     .expect("assert_no_new_utxo failed");
 }
 
-/// Assert that calling update_balance returns a transient error.
-pub async fn assert_temporarily_unavailable(agent: &CkBtcMinterAgent, subaccount: &Subaccount) {
-    let result = agent
-        .update_balance(UpdateBalanceArgs {
-            owner: None,
-            subaccount: Some(*subaccount),
-        })
-        .await
-        .expect("Error while calling update_balance");
-    match result {
-        Ok(utxos_statues) => {
-            for status in utxos_statues {
-                assert_matches!(status, Checked(_));
+/// Assert that calling update_balance returns either a `TemporarilyUnavailable`
+/// error or an `Ok` result where all UTXOs have `Checked` status.
+/// Retries on transient transport errors.
+pub async fn assert_temporarily_unavailable(
+    agent: &CkBtcMinterAgent,
+    logger: &Logger,
+    subaccount: &Subaccount,
+) {
+    ic_system_test_driver::retry_with_msg_async!(
+        "assert_temporarily_unavailable",
+        logger,
+        SHORT_TIMEOUT,
+        RETRY_BACKOFF,
+        || async {
+            let result = match agent
+                .update_balance(UpdateBalanceArgs {
+                    owner: None,
+                    subaccount: Some(*subaccount),
+                })
+                .await
+            {
+                Err(CkBtcMinterAgentError::AgentError(AgentError::TransportError(e))) => {
+                    return Err(anyhow::anyhow!("transport error: {e}"));
+                }
+                other => other.expect("Error while calling update_balance"),
+            };
+            match result {
+                Ok(utxos_statues) => {
+                    for status in utxos_statues {
+                        assert_matches!(status, Checked(_));
+                    }
+                }
+                Err(UpdateBalanceError::TemporarilyUnavailable(..)) => {}
+                Err(error) => {
+                    panic!("Expected TemporarilyUnavailable or Checked but got: {error:?}")
+                }
             }
+            Ok(())
         }
-        Err(error) => {
-            assert_matches!(error, UpdateBalanceError::TemporarilyUnavailable(..));
-        }
-    }
+    )
+    .await
+    .expect("assert_temporarily_unavailable failed");
 }
 
 /// Get transactions log from the ledger canister.