add test to test blockfrommessage functionality

roles/tests-integration/lib/sniffer.rs

@@ -66,34 +66,59 @@ pub struct Sniffer {
 #[derive(Debug, Clone)]
 pub enum Action {
     /// Blocks the message stream after encountering a specific message.
-    BlockFromMessage(BlockFromMessage),
+    IgnoreFromMessage(IgnoreFromMessage),
     /// Intercepts and modifies a message before forwarding it.
     InterceptMessage(Box<InterceptMessage>),
 }
 
+impl Action {
+    /// Returns the action if it is a downstream `IgnoreFromMessage` or `InterceptMessage`  
+    /// with the specified message type.
+    pub fn filter_downstream(
+        action: &Option<Action>,
+        msg_type: MsgType,
+        direction: MessageDirection,
+    ) -> Option<&Self> {
+        action.as_ref().and_then(|action| match action {
+            Action::IgnoreFromMessage(bm)
+                if bm.direction == direction && bm.expected_message_type == msg_type =>
+            {
+                Some(action)
+            }
+
+            Action::InterceptMessage(im)
+                if im.direction == direction && im.expected_message_type == msg_type =>
+            {
+                Some(action)
+            }
+
+            _ => None,
+        })
+    }
+}
 /// Defines an action that blocks the message stream after detecting a specific message.
 #[derive(Debug, Clone)]
-pub struct BlockFromMessage {
+pub struct IgnoreFromMessage {
     direction: MessageDirection,
     expected_message_type: MsgType,
 }
 
-impl BlockFromMessage {
-    /// Creates a new [`BlockFromMessage`] action.
+impl IgnoreFromMessage {
+    /// Creates a new [`IgnoreFromMessage`] action.
     ///
     /// - `direction`: The direction of the message stream to block.
     /// - `expected_message_type`: The type of message after which the stream should be blocked.
     pub fn new(direction: MessageDirection, expected_message_type: MsgType) -> Self {
-        BlockFromMessage {
+        IgnoreFromMessage {
             direction,
             expected_message_type,
         }
     }
 }
 
-impl From<BlockFromMessage> for Action {
-    fn from(value: BlockFromMessage) -> Self {
-        Action::BlockFromMessage(value)
+impl From<IgnoreFromMessage> for Action {
+    fn from(value: IgnoreFromMessage) -> Self {
+        Action::IgnoreFromMessage(value)
     }
 }
 
@@ -283,7 +308,7 @@ impl Sniffer {
             });
             if let Some(ref action) = action {
                 match action {
-                    Action::BlockFromMessage(_) => {
+                    Action::IgnoreFromMessage(_) => {
                         blocked = true;
                         continue;
                     }
@@ -349,7 +374,7 @@ impl Sniffer {
 
             if let Some(ref action) = action {
                 match action {
-                    Action::BlockFromMessage(_) => {
+                    Action::IgnoreFromMessage(_) => {
                         blocked = true;
                         continue;
                     }

roles/tests-integration/tests/sniffer_integration.rs

@@ -10,7 +10,7 @@ use const_sv2::{
     MESSAGE_TYPE_SETUP_CONNECTION, MESSAGE_TYPE_SETUP_CONNECTION_SUCCESS,
     MESSAGE_TYPE_SET_NEW_PREV_HASH,
 };
-use integration_tests_sv2::*;
+use integration_tests_sv2::{sniffer::IgnoreFromMessage, *};
 use roles_logic_sv2::{
     common_messages_sv2::{Protocol, SetupConnection, SetupConnectionError},
     parsers::{CommonMessages, PoolMessages},
@@ -146,40 +146,53 @@ async fn test_sniffer_wait_for_message_type_with_remove() {
 /// Verifies that [`Sniffer`] can intercept and block a message stream.
 ///
 /// This test sets up a chain where a message from the Template Provider (TP)  
-/// passes through two sniffers (`sniffer_a` and `sniffer_b`) before reaching the Pool.
+/// passes through three sniffers (`sniffer_a`, `sniffer_b` and `sniffer_c`) before reaching the
+/// Pool.
 ///
-/// - `sniffer_a` is configured to block `SetupConnectionSuccess` messages directed downstream.
-/// - `sniffer_b` should receive no messages after initial setup, ensuring the block works.
+/// - `sniffer_a` is configured to intercept `SetupConnectionSuccess` messages directed downstream.
+/// - `sniffer_b` is configured to block `SetupConnectionSuccess` messages directed downstream.
+/// - `sniffer_c` should receive no messages after initial setup, ensuring the block works.
 ///
 /// **Flow:**  
-/// `TP -> sniffer_a -> sniffer_b -> Pool`  
+/// `TP -> sniffer_a -> sniffer_b -> sniffer_c -> Pool`  
 #[tokio::test]
 async fn test_sniffer_blocks_message() {
     start_tracing();
     let (_tp, tp_addr) = start_template_provider(None);
 
     // Define an action to block SetupConnectionSuccess messages going downstream.
-    let block_from_message = BlockFromMessage::new(
+    let block_from_message = IgnoreFromMessage::new(
         MessageDirection::ToDownstream,
         MESSAGE_TYPE_SETUP_CONNECTION_SUCCESS,
     );
 
     // `sniffer_a` intercepts and blocks `SetupConnectionSuccess` messages.
-    let (_sniffer_a, sniffer_a_addr) = start_sniffer(
-        "A".to_string(),
-        tp_addr,
+    let (sniffer_a, sniffer_a_addr) = start_sniffer("A".to_string(), tp_addr, false, None).await;
+
+    // `sniffer_b` is placed downstream of `sniffer_a` and should receive nothing.
+    let (_sniffer_b, sniffer_b_addr) = start_sniffer(
+        "B".to_string(),
+        sniffer_a_addr,
         false,
         Some(block_from_message.into()),
     )
     .await;
 
-    // `sniffer_b` is placed downstream of `sniffer_a` and should receive nothing.
-    let (sniffer_b, sniffer_b_addr) =
-        start_sniffer("B".to_string(), sniffer_a_addr, false, None).await;
+    // `sniffer_c` is placed downstream of `sniffer_b` and should receive nothing.
+    let (sniffer_c, sniffer_c_addr) =
+        start_sniffer("C".to_string(), sniffer_b_addr, false, None).await;
 
-    // Start the Pool, connected to `sniffer_b`.
-    let _ = start_pool(Some(sniffer_b_addr)).await;
+    // Start the Pool, connected to `sniffer_c`.
+    let _ = start_pool(Some(sniffer_c_addr)).await;
+
+    // Waiting for intercepting setup connection success on sniffer_a
+    sniffer_a
+        .wait_for_message_type(
+            MessageDirection::ToDownstream,
+            MESSAGE_TYPE_SETUP_CONNECTION_SUCCESS,
+        )
+        .await;
 
-    // Assert that `sniffer_b` does not receive any messages, confirming `sniffer_a`'s block works.
-    assert!(sniffer_b.next_message_from_upstream().is_none());
+    // Assert that `sniffer_c` does not receive any messages, confirming `sniffer_b`'s block works.
+    assert!(sniffer_c.next_message_from_upstream().is_none());
 }