Sniffer::wait_for_message_count

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.
     //