fix: deflake //rs/tests/idx:ii_delegation_test

rs/tests/crypto/canister_sig_verification_cache_test.rs

@@ -115,6 +115,7 @@ pub fn test(env: TestEnv) {
                 signed_delegation: user.signed_delegation.clone(),
                 delegation_identity,
                 polling_timeout: UPDATE_POLLING_TIMEOUT,
+                log: env.logger(),
             })
             .collect();
 
@@ -285,7 +286,8 @@ async fn increment_counter_canister(
     );
     let _ = app_agents_with_delegation[user_i]
         .update(&ctr_canister_id, "write", Blob(vec![]))
-        .await;
+        .await
+        .expect("Update call on counter canister failed");
 }
 
 async fn scrape_metrics_and_check_cache_stats(env: &TestEnv, user_i: usize, call_j: usize) {

rs/tests/driver/src/util/delegations.rs

@@ -1,6 +1,7 @@
 use crate::util::{expiry_time, sign_query, sign_update};
 
 use super::sign_read_state;
+use anyhow::{Context, bail};
 use candid::{CandidType, Deserialize, Principal};
 use canister_test::PrincipalId;
 use ic_agent::Agent;
@@ -16,9 +17,11 @@ use ic_universal_canister::{UNIVERSAL_CANISTER_WASM, wasm};
 use ic_utils::interfaces::ManagementCanister;
 use reqwest::{Client, Response};
 use serde_bytes::ByteBuf;
-use std::time::{Duration, Instant};
+use std::time::Duration;
+
+pub const UPDATE_POLLING_TIMEOUT: Duration = Duration::from_secs(60);
+pub const UPDATE_POLLING_BACKOFF: Duration = Duration::from_millis(500);
 
-pub const UPDATE_POLLING_TIMEOUT: Duration = Duration::from_secs(10);
 /// user ids start with 10000 and increase by 1 for each new user
 pub const USER_NUMBER_OFFSET: u64 = 10_000;
 
@@ -122,6 +125,7 @@ pub struct AgentWithDelegation<'a> {
     pub signed_delegation: SignedDelegation,
     pub delegation_identity: &'a BasicIdentity,
     pub polling_timeout: Duration,
+    pub log: slog::Logger,
 }
 
 impl AgentWithDelegation<'_> {
@@ -130,7 +134,7 @@ impl AgentWithDelegation<'_> {
         method: &str,
         canister_id: &Principal,
         body: Vec<u8>,
-    ) -> Response {
+    ) -> anyhow::Result<Response> {
         let client = Client::new();
         client
             .post(format!(
@@ -143,7 +147,7 @@ impl AgentWithDelegation<'_> {
             .body(body)
             .send()
             .await
-            .unwrap()
+            .context(format!("failed to send {method} request to {canister_id}"))
     }
 
     fn sender(&self) -> Blob {
@@ -154,7 +158,12 @@ impl AgentWithDelegation<'_> {
         )
     }
 
-    pub async fn update(&self, canister_id: &Principal, method_name: &str, arg: Blob) -> MessageId {
+    pub async fn update(
+        &self,
+        canister_id: &Principal,
+        method_name: &str,
+        arg: Blob,
+    ) -> anyhow::Result<MessageId> {
         let update = HttpCanisterUpdate {
             canister_id: Blob(canister_id.as_slice().to_vec()),
             method_name: method_name.to_string(),
@@ -179,16 +188,19 @@ impl AgentWithDelegation<'_> {
             sender_sig: Some(Blob(signature.signature.unwrap())),
         };
         let body = serde_cbor::ser::to_vec(&envelope).unwrap();
-        let _ = self.send_http_request("call", canister_id, body).await;
-        request_id
+        let _ = self
+            .send_http_request("call", canister_id, body)
+            .await
+            .context("failed to send update call")?;
+        Ok(request_id)
     }
 
     pub async fn query(
         &self,
         canister_id: &Principal,
         method_name: &str,
         arg: Blob,
-    ) -> HttpQueryResponse {
+    ) -> anyhow::Result<HttpQueryResponse> {
         let content = HttpQueryContent::Query {
             query: HttpUserQuery {
                 canister_id: Blob(canister_id.as_slice().to_vec()),
@@ -213,9 +225,15 @@ impl AgentWithDelegation<'_> {
             sender_sig: Some(Blob(signature.signature.unwrap())),
         };
         let body = serde_cbor::ser::to_vec(&envelope).unwrap();
-        let response = self.send_http_request("query", canister_id, body).await;
-        let response_bytes = response.bytes().await.unwrap();
-        serde_cbor::from_slice(&response_bytes).unwrap()
+        let response = self
+            .send_http_request("query", canister_id, body)
+            .await
+            .context("failed to send query request")?;
+        let response_bytes = response
+            .bytes()
+            .await
+            .context("failed to read query response bytes")?;
+        serde_cbor::from_slice(&response_bytes).context("failed to deserialize query response")
     }
 
     pub async fn update_and_wait(
@@ -224,7 +242,10 @@ impl AgentWithDelegation<'_> {
         method_name: &str,
         arg: Blob,
     ) -> Result<Vec<u8>, String> {
-        let request_id = self.update(canister_id, method_name, arg.clone()).await;
+        let request_id = self
+            .update(canister_id, method_name, arg.clone())
+            .await
+            .map_err(|e| e.to_string())?;
         let content = HttpReadStateContent::ReadState {
             read_state: HttpReadState {
                 sender: self.sender(),
@@ -250,45 +271,90 @@ impl AgentWithDelegation<'_> {
             )]),
         };
         let body = serde_cbor::ser::to_vec(&read_state_envelope).unwrap();
-        let path = {
-            let p1: &[u8] = b"request_status";
-            let p2: &[u8] = request_id.as_bytes();
-            let p3: &[u8] = b"reply";
-            vec![p1, p2, p3]
-        };
-        let start = Instant::now();
-        let read_state: Vec<u8> = loop {
-            if start.elapsed() > self.polling_timeout {
-                return Err(format!(
-                    "Polling timeout of {} ms was reached",
-                    self.polling_timeout.as_millis()
-                ));
-            }
-            let response = self
-                .send_http_request("read_state", canister_id, body.clone())
-                .await;
-            let read_state_body = response.bytes().await.unwrap();
-            let response_bytes: HttpReadStateResponse =
-                serde_cbor::from_slice(&read_state_body).unwrap();
-            let certificate: Certificate =
-                serde_cbor::from_slice(&response_bytes.certificate).unwrap();
-            let lookup_status = certificate.tree.lookup(&path);
-            match lookup_status {
-                ic_crypto_tree_hash::LookupStatus::Found(x) => match x {
-                    ic_crypto_tree_hash::MixedHashTree::Leaf(y) => break y.clone(),
-                    _ => panic!("Unexpected result from the read_state tree hash structure"),
-                },
-                ic_crypto_tree_hash::LookupStatus::Absent => {
-                    // If request is absent, keep polling
-                    continue;
-                }
-                ic_crypto_tree_hash::LookupStatus::Unknown => {
-                    // If request is unknown, keep polling
-                    continue;
+        let status_path: Vec<&[u8]> = vec![b"request_status", request_id.as_bytes(), b"status"];
+        let reply_path: Vec<&[u8]> = vec![b"request_status", request_id.as_bytes(), b"reply"];
+        let reject_message_path: Vec<&[u8]> =
+            vec![b"request_status", request_id.as_bytes(), b"reject_message"];
+        // The closure returns anyhow::Result<Result<Vec<u8>, String>>:
+        //   - bail!(...) → triggers retry (non-terminal)
+        //   - Ok(Err(...)) → terminal failure, stops retrying immediately
+        //   - Ok(Ok(...)) → success
+        let read_state_result: Result<Vec<u8>, String> = crate::retry_with_msg_async!(
+            "update_and_wait read_state polling",
+            &self.log,
+            self.polling_timeout,
+            UPDATE_POLLING_BACKOFF,
+            || async {
+                let response = self
+                    .send_http_request("read_state", canister_id, body.clone())
+                    .await
+                    .context("failed to send read_state request")?;
+                let read_state_body = response.bytes().await.context("failed to read read_state response bytes")?;
+                let response_bytes: HttpReadStateResponse =
+                    serde_cbor::from_slice(&read_state_body).context("failed to deserialize HttpReadStateResponse")?;
+                let certificate: Certificate =
+                    serde_cbor::from_slice(&response_bytes.certificate).context("failed to deserialize Certificate")?;
+                let tree = &certificate.tree;
+
+                // First, check the status path to determine the request state.
+                let status_str = match tree.lookup(&status_path) {
+                    ic_crypto_tree_hash::LookupStatus::Found(
+                        ic_crypto_tree_hash::MixedHashTree::Leaf(s),
+                    ) => String::from_utf8(s.clone()).context("request status should be valid utf8")?,
+                    ic_crypto_tree_hash::LookupStatus::Absent
+                    | ic_crypto_tree_hash::LookupStatus::Unknown => {
+                        bail!("Request status not yet available, keep polling")
+                    }
+                    ic_crypto_tree_hash::LookupStatus::Found(_) => {
+                        return Ok(Err(format!(
+                            "Update call to {canister_id}/{method_name} (request_id={request_id}): \
+                             unexpected non-leaf node at request status path"
+                        )));
+                    }
+                };
+
+                match status_str.as_str() {
+                    "replied" => {
+                        // Read the reply payload.
+                        match tree.lookup(&reply_path) {
+                            ic_crypto_tree_hash::LookupStatus::Found(
+                                ic_crypto_tree_hash::MixedHashTree::Leaf(y),
+                            ) => Ok(Ok(y.clone())),
+                            _ => Ok(Err(format!(
+                                "Update call to {canister_id}/{method_name} (request_id={request_id}): \
+                                 status is 'replied' but reply leaf is missing or malformed"
+                            ))),
+                        }
+                    }
+                    "rejected" => {
+                        let message = match tree.lookup(&reject_message_path) {
+                            ic_crypto_tree_hash::LookupStatus::Found(
+                                ic_crypto_tree_hash::MixedHashTree::Leaf(m),
+                            ) => String::from_utf8_lossy(m).to_string(),
+                            _ => "no message".to_string(),
+                        };
+                        Ok(Err(format!(
+                            "Update call to {canister_id}/{method_name} (request_id={request_id}) \
+                             was rejected with message: {message}"
+                        )))
+                    }
+                    "done" => Ok(Err(format!(
+                        "Update call to {canister_id}/{method_name} (request_id={request_id}) \
+                         reached terminal state 'done' without a reply"
+                    ))),
+                    "pruned" => Ok(Err(format!(
+                        "Update call to {canister_id}/{method_name} (request_id={request_id}) \
+                         reached terminal state 'pruned': reply/reject data has been pruned"
+                    ))),
+                    _ => {
+                        bail!("Request status is '{status_str}', keep polling")
+                    }
                 }
-            };
-        };
-        Ok(read_state)
+            }
+        )
+        .await
+        .map_err(|e| e.to_string())?;
+        read_state_result
     }
 }
 

rs/tests/idx/ii_delegation_test.rs

@@ -10,14 +10,11 @@ use ic_system_test_driver::driver::test_env_api::{
 };
 use ic_system_test_driver::systest;
 use ic_system_test_driver::util::delegations::*;
-use ic_system_test_driver::util::{
-    agent_with_identity, assert_canister_counter_with_retries, block_on, random_ed25519_identity,
-};
+use ic_system_test_driver::util::{agent_with_identity, block_on, random_ed25519_identity};
 use ic_types::messages::{Blob, HttpQueryResponse};
 use ic_universal_canister::wasm;
 use slog::info;
 use std::env;
-use std::time::Duration;
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
@@ -106,37 +103,42 @@ pub fn test(env: TestEnv) {
         signed_delegation,
         delegation_identity: &delegation_identity,
         polling_timeout: UPDATE_POLLING_TIMEOUT,
+        log: log.clone(),
     };
     info!(
         log,
         "Making an update call on counter canister with delegation (increment counter)"
     );
-    let _ = block_on(app_agent_with_delegation.update(&counter_canister_id, "write", Blob(vec![])));
+    block_on(app_agent_with_delegation.update_and_wait(
+        &counter_canister_id,
+        "write",
+        Blob(vec![]),
+    ))
+    .expect("Update call on counter canister failed");
     info!(
         log,
         "Making a query call on counter canister with delegation (read counter)"
     );
     let query_response =
-        block_on(app_agent_with_delegation.query(&counter_canister_id, "read", Blob(vec![])));
+        block_on(app_agent_with_delegation.query(&counter_canister_id, "read", Blob(vec![])))
+            .expect("Query call on counter canister failed");
     match query_response {
-        HttpQueryResponse::Replied { .. } => (),
+        HttpQueryResponse::Replied { reply } => {
+            let counter = u32::from_le_bytes(
+                reply
+                    .arg
+                    .as_slice()
+                    .try_into()
+                    .expect("slice with incorrect length"),
+            );
+            assert_eq!(counter, 1, "Counter canister should have value 1");
+        }
         HttpQueryResponse::Rejected {
             error_code,
             reject_message,
             ..
         } => panic!("Query call was rejected: code={error_code}, message={reject_message}"),
     }
-    info!(log, "Asserting canister counter has value=1");
-    let app_agent = app_node.build_default_agent();
-    block_on(assert_canister_counter_with_retries(
-        &log,
-        &app_agent,
-        &counter_canister_id,
-        vec![],
-        1,
-        10,
-        Duration::from_secs(1),
-    ));
     let expected_principal = Principal::self_authenticating(&ii_derived_public_key);
     info!(
         log,
@@ -148,7 +150,8 @@ pub fn test(env: TestEnv) {
             &ucan_id,
             "query",
             Blob(wasm().caller().append_and_reply().build()),
-        ));
+        ))
+        .expect("Query call on ucan canister failed");
         match response {
             HttpQueryResponse::Replied { reply } => Principal::from_slice(reply.arg.as_ref()),
             HttpQueryResponse::Rejected { reject_message, .. } => {