From 27eec57a338ab7750c720de1c5d10265f5243ec8 Mon Sep 17 00:00:00 2001
From: jbesraa <jbesraa@gmail.com>
Date: Mon, 18 Nov 2024 11:17:57 +0200
Subject: [PATCH] Add `check_on_drop` property to `Sniffer`

..If this flag is set to true, the `Sniffer` **will**
validate that all the downstrea and upstream messages were handled
before dropping, otherwise they are ignored.
---
 roles/tests-integration/tests/common/mod.rs   |  3 +-
 .../tests-integration/tests/common/sniffer.rs | 39 +++++++++++--------
 .../tests/pool_integration.rs                 |  9 ++++-
 3 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/roles/tests-integration/tests/common/mod.rs b/roles/tests-integration/tests/common/mod.rs
index 10622e3d3..f05a0d272 100644
--- a/roles/tests-integration/tests/common/mod.rs
+++ b/roles/tests-integration/tests/common/mod.rs
@@ -198,8 +198,9 @@ 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).await;
+    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;
diff --git a/roles/tests-integration/tests/common/sniffer.rs b/roles/tests-integration/tests/common/sniffer.rs
index be1d68f25..9c6cd112b 100644
--- a/roles/tests-integration/tests/common/sniffer.rs
+++ b/roles/tests-integration/tests/common/sniffer.rs
@@ -51,6 +51,7 @@ pub struct Sniffer {
     upstream_address: SocketAddr,
     messages_from_downstream: MessagesAggregator,
     messages_from_upstream: MessagesAggregator,
+    check_on_drop: bool,
 }
 
 impl Sniffer {
@@ -60,6 +61,7 @@ impl Sniffer {
         identifier: String,
         listening_address: SocketAddr,
         upstream_address: SocketAddr,
+        check_on_drop: bool,
     ) -> Self {
         Self {
             identifier,
@@ -67,6 +69,7 @@ impl Sniffer {
             upstream_address,
             messages_from_downstream: MessagesAggregator::new(),
             messages_from_upstream: MessagesAggregator::new(),
+            check_on_drop,
         }
     }
 
@@ -436,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!(
-                "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!();
+        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!();
+            }
         }
     }
 }
diff --git a/roles/tests-integration/tests/pool_integration.rs b/roles/tests-integration/tests/pool_integration.rs
index 2f8430adb..942ca1c50 100644
--- a/roles/tests-integration/tests/pool_integration.rs
+++ b/roles/tests-integration/tests/pool_integration.rs
@@ -17,7 +17,14 @@ async fn success_pool_template_provider_connection() {
     let _tp = common::start_template_provider(tp_addr.port()).await;
     let sniffer_identifier =
         "success_pool_template_provider_connection tp_pool sniffer".to_string();
-    let sniffer = common::start_sniffer(sniffer_identifier, sniffer_addr, tp_addr).await;
+    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