From bdc28b8d7fdf0f49e3443acf1bc8c4e833930411 Mon Sep 17 00:00:00 2001
From: gterzian <2792687+gterzian@users.noreply.github.com>
Date: Thu, 29 Sep 2022 17:11:10 +0800
Subject: [PATCH] protocol: adding test for operation re-ask

---
 .../src/tests/operations_scenarios.rs         | 151 ++++++++++++++++++
 1 file changed, 151 insertions(+)

diff --git a/massa-protocol-worker/src/tests/operations_scenarios.rs b/massa-protocol-worker/src/tests/operations_scenarios.rs
index 7cc31de114e..684354594fc 100644
--- a/massa-protocol-worker/src/tests/operations_scenarios.rs
+++ b/massa-protocol-worker/src/tests/operations_scenarios.rs
@@ -755,6 +755,157 @@ async fn test_protocol_ask_operations_on_batch_received() {
     .await;
 }
 
+#[tokio::test]
+#[serial]
+async fn test_protocol_re_ask_operations_to_another_node_on_batch_received_after_delay() {
+    let protocol_config = &tools::PROTOCOL_CONFIG;
+    protocol_test(
+        protocol_config,
+        async move |mut network_controller,
+                    protocol_event_receiver,
+                    protocol_command_sender,
+                    protocol_manager,
+                    protocol_pool_event_receiver| {
+            // Create 3 node.
+            let mut nodes = tools::create_and_connect_nodes(3, &mut network_controller).await;
+
+            let creator_node = nodes.pop().expect("Failed to get node info.");
+
+            // 1. Create an operation
+            let operation = tools::create_operation_with_expire_period(&creator_node.keypair, 1);
+
+            let expected_operation_id = operation.id;
+            // 3. Send operation batch to protocol.
+            network_controller
+                .send_operation_batch(creator_node.id, vec![expected_operation_id])
+                .await;
+
+            // First ask
+            let first_asked_to_node;
+            match network_controller
+                .wait_command(1000.into(), |cmd| match cmd {
+                    cmd @ NetworkCommand::AskForOperations { .. } => Some(cmd),
+                    _ => None,
+                })
+                .await
+            {
+                Some(NetworkCommand::AskForOperations { to_node, wishlist }) => {
+                    assert_eq!(wishlist.len(), 1);
+                    assert!(wishlist.contains(&expected_operation_id.prefix()));
+                    first_asked_to_node = to_node;
+                }
+                _ => panic!("Unexpected or no network command."),
+            };
+
+            // Second announcement from other node.
+            network_controller
+                .send_operation_batch(nodes[1].id, vec![expected_operation_id])
+                .await;
+
+            // Second ask, to a different node.
+            match network_controller
+                .wait_command(1000.into(), |cmd| match cmd {
+                    cmd @ NetworkCommand::AskForOperations { .. } => Some(cmd),
+                    _ => None,
+                })
+                .await
+            {
+                Some(NetworkCommand::AskForOperations { to_node, wishlist }) => {
+                    assert_eq!(wishlist.len(), 1);
+                    assert!(wishlist.contains(&expected_operation_id.prefix()));
+                    assert_ne!(to_node, first_asked_to_node);
+                }
+                _ => panic!("Unexpected or no network command."),
+            };
+            (
+                network_controller,
+                protocol_event_receiver,
+                protocol_command_sender,
+                protocol_manager,
+                protocol_pool_event_receiver,
+            )
+        },
+    )
+    .await;
+}
+
+#[tokio::test]
+#[serial]
+async fn test_protocol_does_not_re_ask_operations_to_another_node_if_received() {
+    let protocol_config = &tools::PROTOCOL_CONFIG;
+    protocol_test(
+        protocol_config,
+        async move |mut network_controller,
+                    protocol_event_receiver,
+                    protocol_command_sender,
+                    protocol_manager,
+                    protocol_pool_event_receiver| {
+            // Create 3 node.
+            let mut nodes = tools::create_and_connect_nodes(3, &mut network_controller).await;
+
+            let creator_node = nodes.pop().expect("Failed to get node info.");
+
+            // 1. Create an operation
+            let operation = tools::create_operation_with_expire_period(&creator_node.keypair, 1);
+
+            let expected_operation_id = operation.id;
+            // 3. Send operation batch to protocol.
+            network_controller
+                .send_operation_batch(creator_node.id, vec![expected_operation_id])
+                .await;
+
+            // First ask
+            match network_controller
+                .wait_command(1000.into(), |cmd| match cmd {
+                    cmd @ NetworkCommand::AskForOperations { .. } => Some(cmd),
+                    _ => None,
+                })
+                .await
+            {
+                Some(NetworkCommand::AskForOperations { to_node, wishlist }) => {
+                    assert_eq!(wishlist.len(), 1);
+                    assert!(wishlist.contains(&expected_operation_id.prefix()));
+                    assert_eq!(to_node, creator_node.id);
+                }
+                _ => panic!("Unexpected or no network command."),
+            };
+
+            // Send operation to protocol.
+            network_controller
+                .send_operations(creator_node.id, vec![operation])
+                .await;
+
+            // Second announcement from other node.
+            network_controller
+                .send_operation_batch(nodes[1].id, vec![expected_operation_id])
+                .await;
+
+            // Second ask, to a different node, should not occur,
+            // because the operation has been received in the meantime.
+            match network_controller
+                .wait_command(1000.into(), |cmd| match cmd {
+                    cmd @ NetworkCommand::AskForOperations { .. } => Some(cmd),
+                    _ => None,
+                })
+                .await
+            {
+                Some(NetworkCommand::AskForOperations { .. }) => {
+                    panic!("Unexpected ask for operations");
+                }
+                _ => {}
+            };
+            (
+                network_controller,
+                protocol_event_receiver,
+                protocol_command_sender,
+                protocol_manager,
+                protocol_pool_event_receiver,
+            )
+        },
+    )
+    .await;
+}
+
 #[tokio::test]
 #[serial]
 async fn test_protocol_on_ask_operations() {