From cf01c07b86c474c6518b4f47e5a4e55ec528a00d Mon Sep 17 00:00:00 2001
From: plebhash <plebhash@proton.me>
Date: Tue, 10 Dec 2024 13:25:09 +0700
Subject: [PATCH] Sniffer::wait_for_message_count

---
 .../tests-integration/tests/common/sniffer.rs | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/roles/tests-integration/tests/common/sniffer.rs b/roles/tests-integration/tests/common/sniffer.rs
index 985f1ba92..0e61fb4c9 100644
--- a/roles/tests-integration/tests/common/sniffer.rs
+++ b/roles/tests-integration/tests/common/sniffer.rs
@@ -418,6 +418,25 @@ impl Sniffer {
             panic!("Impossible to accept dowsntream connection")
         }
     }
+
+    /// used to block the test runtime
+    /// while we wait until Sniffer has a specific count of messages
+    pub async fn wait_for_message_count(&self, message_direction: MessageDirection, count: usize) {
+        loop {
+            let message_count = match message_direction {
+                MessageDirection::ToDownstream => {
+                    self.messages_from_upstream.message_count()
+                },
+                MessageDirection::ToUpstream => {
+                    self.messages_from_downstream.message_count()
+                }
+            };
+
+            if message_count == count {
+                return
+            }
+        }
+    }
 }
 
 // Utility macro to assert that the downstream and upstream roles have sent specific messages.
@@ -589,6 +608,12 @@ impl MessagesAggregator {
             .unwrap()
     }
 
+    fn message_count(&self) -> usize {
+        self.messages
+            .safe_lock(|messages| messages.len())
+            .unwrap()
+    }
+
     // The aggregator queues messages in FIFO order, so this function returns the oldest message in
     // the queue.
     //