Merge pull request #1258 from jbesraa/2024-11-18-sniffer-improvements

`Sniffer` Improvements

roles/tests-integration/tests/common/mod.rs

@@ -194,8 +194,13 @@ pub fn get_available_address() -> SocketAddr {
     SocketAddr::from(([127, 0, 0, 1], port))
 }
 
-pub async fn start_sniffer(listening_address: SocketAddr, upstream: SocketAddr) -> Sniffer {
-    let sniffer = Sniffer::new(listening_address, upstream).await;
+pub async fn start_sniffer(
+    identifier: String,
+    listening_address: SocketAddr,
+    upstream: SocketAddr,
+    check_on_drop: bool,
+) -> Sniffer {
+    let sniffer = Sniffer::new(identifier, listening_address, upstream, check_on_drop).await;
     let sniffer_clone = sniffer.clone();
     tokio::spawn(async move {
         sniffer_clone.start().await;

roles/tests-integration/tests/common/sniffer.rs

@@ -46,21 +46,30 @@ enum SnifferError {
 /// messages in a specific order and to inspect the messages details.
 #[derive(Debug, Clone)]
 pub struct Sniffer {
+    identifier: String,
     listening_address: SocketAddr,
     upstream_address: SocketAddr,
     messages_from_downstream: MessagesAggregator,
     messages_from_upstream: MessagesAggregator,
+    check_on_drop: bool,
 }
 
 impl Sniffer {
     /// Creates a new sniffer that listens on the given listening address and connects to the given
     /// upstream address.
-    pub async fn new(listening_address: SocketAddr, upstream_address: SocketAddr) -> Self {
+    pub async fn new(
+        identifier: String,
+        listening_address: SocketAddr,
+        upstream_address: SocketAddr,
+        check_on_drop: bool,
+    ) -> Self {
         Self {
+            identifier,
             listening_address,
             upstream_address,
             messages_from_downstream: MessagesAggregator::new(),
             messages_from_upstream: MessagesAggregator::new(),
+            check_on_drop,
         }
     }
 
@@ -430,23 +439,25 @@ macro_rules! assert_jd_message {
 // This is useful to ensure that the test has checked all exchanged messages between the roles.
 impl Drop for Sniffer {
     fn drop(&mut self) {
-        // Don't print backtrace on panic
-        std::panic::set_hook(Box::new(|_| {
-            println!();
-        }));
-        if !self.messages_from_downstream.is_empty() {
-            println!(
-                "You didn't handle all downstream messages: {:?}",
-                self.messages_from_downstream
-            );
-            panic!();
-        }
-        if !self.messages_from_upstream.is_empty() {
-            println!(
-                "You didn't handle all upstream messages: {:?}",
-                self.messages_from_upstream
-            );
-            panic!();
+        if self.check_on_drop {
+            // Don't print backtrace on panic
+            std::panic::set_hook(Box::new(|_| {
+                println!();
+            }));
+            if !self.messages_from_downstream.is_empty() {
+                println!(
+                    "Sniffer {}: You didn't handle all downstream messages: {:?}",
+                    self.identifier, self.messages_from_downstream
+                );
+                panic!();
+            }
+            if !self.messages_from_upstream.is_empty() {
+                println!(
+                    "Sniffer{}: You didn't handle all upstream messages: {:?}",
+                    self.identifier, self.messages_from_upstream
+                );
+                panic!();
+            }
         }
     }
 }

roles/tests-integration/tests/pool_integration.rs

@@ -15,7 +15,16 @@ async fn success_pool_template_provider_connection() {
     let tp_addr = common::get_available_address();
     let pool_addr = common::get_available_address();
     let _tp = common::start_template_provider(tp_addr.port()).await;
-    let sniffer = common::start_sniffer(sniffer_addr, tp_addr).await;
+    let sniffer_identifier =
+        "success_pool_template_provider_connection tp_pool sniffer".to_string();
+    let sniffer_check_on_drop = true;
+    let sniffer = common::start_sniffer(
+        sniffer_identifier,
+        sniffer_addr,
+        tp_addr,
+        sniffer_check_on_drop,
+    )
+    .await;
     let _ = common::start_pool(Some(pool_addr), Some(sniffer_addr)).await;
     // here we assert that the downstream(pool in this case) have sent `SetupConnection` message
     // with the correct parameters, protocol, flags, min_version and max_version.  Note that the