From c205c0b3ab4427260360cee816bd3ab2ae2d646e Mon Sep 17 00:00:00 2001
From: Gustavo Inacio <gustavo@semiotic.ai>
Date: Tue, 11 Feb 2025 20:04:25 +0100
Subject: [PATCH] test: fix nextest flaky test

Signed-off-by: Gustavo Inacio <gustavo@semiotic.ai>
---
 ...469cd7bc10a39e7c94f2b3d57946e4a060ff.json} |  4 +--
 .../tests/sender_account_manager_test.rs      | 34 +++++++++++++------
 2 files changed, 25 insertions(+), 13 deletions(-)
 rename .sqlx/{query-661b961e286d8591d1eaaf29d3bdf5748c361f3d4c4136b8d5b3a3c3af5fcc9e.json => query-3d76b9e04b0c09999c5653e4176e469cd7bc10a39e7c94f2b3d57946e4a060ff.json} (92%)

diff --git a/.sqlx/query-661b961e286d8591d1eaaf29d3bdf5748c361f3d4c4136b8d5b3a3c3af5fcc9e.json b/.sqlx/query-3d76b9e04b0c09999c5653e4176e469cd7bc10a39e7c94f2b3d57946e4a060ff.json
similarity index 92%
rename from .sqlx/query-661b961e286d8591d1eaaf29d3bdf5748c361f3d4c4136b8d5b3a3c3af5fcc9e.json
rename to .sqlx/query-3d76b9e04b0c09999c5653e4176e469cd7bc10a39e7c94f2b3d57946e4a060ff.json
index 65352b480..aecfd0e24 100644
--- a/.sqlx/query-661b961e286d8591d1eaaf29d3bdf5748c361f3d4c4136b8d5b3a3c3af5fcc9e.json
+++ b/.sqlx/query-3d76b9e04b0c09999c5653e4176e469cd7bc10a39e7c94f2b3d57946e4a060ff.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n            SELECT * FROM scalar_tap_ravs WHERE last = true;\n        ",
+  "query": "\n            SELECT * FROM scalar_tap_ravs WHERE last;\n        ",
   "describe": {
     "columns": [
       {
@@ -64,5 +64,5 @@
       true
     ]
   },
-  "hash": "661b961e286d8591d1eaaf29d3bdf5748c361f3d4c4136b8d5b3a3c3af5fcc9e"
+  "hash": "3d76b9e04b0c09999c5653e4176e469cd7bc10a39e7c94f2b3d57946e4a060ff"
 }
diff --git a/crates/tap-agent/tests/sender_account_manager_test.rs b/crates/tap-agent/tests/sender_account_manager_test.rs
index ed1dcbd42..65fa85bf7 100644
--- a/crates/tap-agent/tests/sender_account_manager_test.rs
+++ b/crates/tap-agent/tests/sender_account_manager_test.rs
@@ -13,10 +13,10 @@ use indexer_tap_agent::{
         create_received_receipt, create_sender_accounts_manager, store_receipt, ALLOCATION_ID_0,
     },
 };
-use ractor::ActorRef;
+use ractor::{ActorRef, ActorStatus};
 use serde_json::json;
 use sqlx::PgPool;
-use test_assets::{flush_messages, TAP_SENDER as SENDER, TAP_SIGNER as SIGNER};
+use test_assets::{assert_while_retry, flush_messages, TAP_SENDER as SENDER, TAP_SIGNER as SIGNER};
 use wiremock::{
     matchers::{body_string_contains, method},
     Mock, MockServer, ResponseTemplate,
@@ -74,9 +74,9 @@ async fn sender_account_manager_layer_test(pgpool: PgPool) {
     flush_messages(&notify).await;
 
     // verify if create sender account
-    let actor_ref =
+    let sender_account_ref =
         ActorRef::<SenderAccountMessage>::where_is(format!("{}:{}", prefix.clone(), SENDER.1));
-    assert!(actor_ref.is_some());
+    assert!(sender_account_ref.is_some());
 
     let receipt = create_received_receipt(&ALLOCATION_ID_0, &SIGNER.0, 1, 1, TRIGGER_VALUE - 10);
     store_receipt(&pgpool, receipt.signed_receipt())
@@ -84,7 +84,7 @@ async fn sender_account_manager_layer_test(pgpool: PgPool) {
         .unwrap();
 
     // we expect it to create a sender allocation
-    actor_ref
+    sender_account_ref
         .clone()
         .unwrap()
         .cast(SenderAccountMessage::UpdateAllocationIds(
@@ -93,15 +93,28 @@ async fn sender_account_manager_layer_test(pgpool: PgPool) {
                 .collect(),
         ))
         .unwrap();
-    flush_messages(&notify).await;
+
+    assert_while_retry!({
+        ActorRef::<SenderAllocationMessage>::where_is(format!(
+            "{}:{}:{}",
+            prefix, SENDER.1, ALLOCATION_ID_0,
+        ))
+        .is_none()
+    });
+    let allocation_ref = ActorRef::<SenderAllocationMessage>::where_is(format!(
+        "{}:{}:{}",
+        prefix, SENDER.1, ALLOCATION_ID_0,
+    ))
+    .unwrap();
 
     // try to delete sender allocation_id
-    actor_ref
+    sender_account_ref
         .clone()
         .unwrap()
         .cast(SenderAccountMessage::UpdateAllocationIds(HashSet::new()))
         .unwrap();
-    flush_messages(&notify).await;
+    allocation_ref.wait(None).await.unwrap();
+    assert_eq!(allocation_ref.get_status(), ActorStatus::Stopped);
 
     assert!(ActorRef::<SenderAllocationMessage>::where_is(format!(
         "{}:{}:{}",
@@ -116,15 +129,14 @@ async fn sender_account_manager_layer_test(pgpool: PgPool) {
         ))
         .unwrap();
 
-    flush_messages(&notify).await;
+    sender_account_ref.unwrap().wait(None).await.unwrap();
     // verify if it gets removed
     let actor_ref = ActorRef::<SenderAccountMessage>::where_is(format!("{}:{}", prefix, SENDER.1));
     assert!(actor_ref.is_none());
 
-    //verify the rav is marked as last
     let rav_marked_as_last = sqlx::query!(
         r#"
-            SELECT * FROM scalar_tap_ravs WHERE last = true;
+            SELECT * FROM scalar_tap_ravs WHERE last;
         "#,
     )
     .fetch_all(&pgpool)