diff --git a/rs/tests/ckbtc/ckbtc_minter_update_balance.rs b/rs/tests/ckbtc/ckbtc_minter_update_balance.rs index 023ef2817f93..9a0165df4f57 100644 --- a/rs/tests/ckbtc/ckbtc_minter_update_balance.rs +++ b/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; diff --git a/rs/tests/ckbtc/src/utils.rs b/rs/tests/ckbtc/src/utils.rs index 2cc71b369f79..9793705e318e 100644 --- a/rs/tests/ckbtc/src/utils.rs +++ b/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.