Updating unit test

chitchat/src/lib.rs

@@ -585,9 +585,23 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn test_dead_node_should_not_be_gossiped_when_node_joins() -> anyhow::Result<()> {
+    async fn test_dead_node_kvs_are_when_node_joins() -> anyhow::Result<()> {
         let transport = ChannelTransport::with_mtu(MAX_UDP_DATAGRAM_PAYLOAD_SIZE);
-        let mut nodes = setup_nodes(40001..=40004, &transport).await;
+        // starting 2 nodes.
+        let mut nodes = setup_nodes(40001..=40002, &transport).await;
+
+        // Let's add a key to node1.
+        let node1_id = {
+            let node1 = nodes.get(0).unwrap();
+            let node1_chitchat = node1.chitchat();
+            node1_chitchat
+                .lock()
+                .await
+                .self_node_state()
+                .set("test_key", "test_val");
+            node1.chitchat_id().clone()
+        };
+
         {
             let node2 = nodes.get(1).unwrap();
             assert_eq!(node2.chitchat_id().advertise_port(), 40002);
@@ -596,54 +610,67 @@ mod tests {
                 &[
                     ChitchatId::for_local_test(40001),
                     ChitchatId::for_local_test(40002),
-                    ChitchatId::for_local_test(40003),
-                    ChitchatId::for_local_test(40004),
                 ],
             )
             .await;
+            let node2_chitchat = node2.chitchat();
+            // We have received node3's key
+            let value = node2_chitchat
+                .lock()
+                .await
+                .node_state(&node1_id)
+                .unwrap()
+                .get("test_key")
+                .unwrap()
+                .to_string();
+            assert_eq!(&value, "test_val");
         }
 
-        // Take node 3 down.
-        let node3 = nodes.remove(2);
-        assert_eq!(node3.chitchat_id().advertise_port(), 40003);
-        node3.shutdown().await.unwrap();
+        // Take node 1 down.
+        let node1 = nodes.remove(0);
+        assert_eq!(node1.chitchat_id().advertise_port(), 40001);
+        node1.shutdown().await.unwrap();
 
+        // Node 2 has detected that node 1 is missing.
         {
-            let node2 = nodes.get(1).unwrap();
+            let node2 = nodes.get(0).unwrap();
             assert_eq!(node2.chitchat_id().advertise_port(), 40002);
-            wait_for_chitchat_state(
-                node2.chitchat(),
-                &[
-                    ChitchatId::for_local_test(40_001),
-                    ChitchatId::for_local_test(40002),
-                    ChitchatId::for_local_test(40_004),
-                ],
-            )
-            .await;
+            wait_for_chitchat_state(node2.chitchat(), &[ChitchatId::for_local_test(40_002)]).await;
         }
 
-        // Restart node at localhost:40003 with new name
-        let mut new_config = ChitchatConfig::for_test(40_003);
+        // Restart node at localhost:40001 with new name
+        let mut new_config = ChitchatConfig::for_test(40_001);
         new_config.chitchat_id.node_id = "new_node".to_string();
         let new_chitchat_id = new_config.chitchat_id.clone();
-        let seed_addr = ChitchatId::for_local_test(40_002).gossip_advertise_addr;
+        let seed_addr = ChitchatId::for_local_test(40_001).gossip_advertise_addr;
         new_config.seed_nodes = vec![seed_addr.to_string()];
-        let new_node_chitchat = spawn_chitchat(new_config, Vec::new(), &transport)
+        let new_node_chitchat_handle = spawn_chitchat(new_config, Vec::new(), &transport)
             .await
             .unwrap();
 
+        let new_node_chitchat = new_node_chitchat_handle.chitchat();
         wait_for_chitchat_state(
-            new_node_chitchat.chitchat(),
-            &[
-                ChitchatId::for_local_test(40_001),
-                ChitchatId::for_local_test(40_002),
-                new_chitchat_id,
-                ChitchatId::for_local_test(40_004),
-            ],
+            new_node_chitchat.clone(),
+            &[ChitchatId::for_local_test(40_002), new_chitchat_id],
         )
         .await;
 
-        nodes.push(new_node_chitchat);
+        {
+            let new_node_chitchat_guard = new_node_chitchat.lock().await;
+            let test_val = new_node_chitchat_guard
+                .node_state(&node1_id)
+                .unwrap()
+                .get("test_key")
+                .unwrap();
+            assert_eq!(test_val, "test_val");
+
+            // Let's check that node1 is seen as dead.
+            let dead_nodes: HashSet<&ChitchatId> = new_node_chitchat_guard.dead_nodes().collect();
+            assert_eq!(dead_nodes.len(), 1);
+            assert!(dead_nodes.contains(&node1_id));
+        }
+
+        nodes.push(new_node_chitchat_handle);
         shutdown_nodes(nodes).await?;
         Ok(())
     }