test: Add network notification testing apparatus

ipvm-wg · Jan 29, 2024 · d3d4c04 · d3d4c04
1 parent 31f697c
commit d3d4c04
Showing 1 changed file with 55 additions and 29 deletions.
diff --git a/homestar-runtime/src/event_handler/notification/swarm.rs b/homestar-runtime/src/event_handler/notification/swarm.rs
@@ -297,51 +297,77 @@ impl TryFrom<Ipld> for ConnectionClosed {
 #[cfg(test)]
 mod test {
     use super::*;
-    use libp2p::PeerId;
+
+    fn generate_notifications(
+        peer_id: PeerId,
+        address: Multiaddr,
+    ) -> Vec<(i64, NetworkNotification)> {
+        let connection_established = ConnectionEstablished::new(peer_id, address.clone());
+        let connection_closed = ConnectionClosed::new(peer_id, address.clone());
+
+        vec![
+            (
+                connection_established.timestamp,
+                NetworkNotification::ConnnectionEstablished(connection_established.clone()),
+            ),
+            (
+                connection_closed.timestamp,
+                NetworkNotification::ConnnectionClosed(connection_closed.clone()),
+            ),
+        ]
+    }
 
     #[test]
     fn notification_bytes_rountrip() {
         let peer_id = PeerId::random();
         let address = Multiaddr::from_str("/ip4/127.0.0.1/tcp/7000").unwrap();
-        let inner = ConnectionEstablished::new(peer_id, address.clone());
-
-        let notification = NetworkNotification::ConnnectionEstablished(inner.clone());
-        let bytes = notification.to_json().unwrap();
-        let parsed = NetworkNotification::from_json(bytes.as_ref()).unwrap();
 
-        match parsed {
-            NetworkNotification::ConnnectionEstablished(n) => {
-                let parsed_peer_id = PeerId::from_str(&n.peer_id).unwrap();
-                let parsed_address = Multiaddr::from_str(&n.address).unwrap();
-
-                assert_eq!(parsed_peer_id, peer_id);
-                assert_eq!(parsed_address, address);
-                assert_eq!(n.timestamp, inner.timestamp);
+        // Generate notifications and convert them to bytes
+        let notifications: Vec<(i64, Vec<u8>)> = generate_notifications(peer_id, address.clone())
+            .into_iter()
+            .map(|(timestamp, notification)| (timestamp, notification.to_json().unwrap()))
+            .collect();
+
+        for (timestamp, bytes) in notifications {
+            match NetworkNotification::from_json(bytes.as_ref()).unwrap() {
+                NetworkNotification::ConnnectionEstablished(n) => {
+                    assert_eq!(PeerId::from_str(&n.peer_id).unwrap(), peer_id);
+                    assert_eq!(Multiaddr::from_str(&n.address).unwrap(), address);
+                    assert_eq!(n.timestamp, timestamp)
+                }
+                NetworkNotification::ConnnectionClosed(n) => {
+                    assert_eq!(PeerId::from_str(&n.peer_id).unwrap(), peer_id);
+                    assert_eq!(Multiaddr::from_str(&n.address).unwrap(), address);
+                    assert_eq!(n.timestamp, timestamp)
+                }
             }
-            _ => panic!("Parsed notification did not matched expected variant"),
         }
     }
 
     #[test]
     fn notification_json_string_rountrip() {
         let peer_id = PeerId::random();
         let address = Multiaddr::from_str("/ip4/127.0.0.1/tcp/7000").unwrap();
-        let inner = ConnectionEstablished::new(peer_id, address.clone());
-
-        let notification = NetworkNotification::ConnnectionEstablished(inner.clone());
-        let json_string = notification.to_json_string().unwrap();
-        let parsed = NetworkNotification::from_json_string(json_string).unwrap();
-
-        match parsed {
-            NetworkNotification::ConnnectionEstablished(n) => {
-                let parsed_peer_id = PeerId::from_str(&n.peer_id).unwrap();
-                let parsed_address = Multiaddr::from_str(&n.address).unwrap();
 
-                assert_eq!(parsed_peer_id, peer_id);
-                assert_eq!(parsed_address, address);
-                assert_eq!(n.timestamp, inner.timestamp);
+        // Generate notifications and convert them to JSON strings
+        let notifications: Vec<(i64, String)> = generate_notifications(peer_id, address.clone())
+            .into_iter()
+            .map(|(timestamp, notification)| (timestamp, notification.to_json_string().unwrap()))
+            .collect();
+
+        for (timestamp, json) in notifications {
+            match NetworkNotification::from_json_string(json).unwrap() {
+                NetworkNotification::ConnnectionEstablished(n) => {
+                    assert_eq!(PeerId::from_str(&n.peer_id).unwrap(), peer_id);
+                    assert_eq!(Multiaddr::from_str(&n.address).unwrap(), address);
+                    assert_eq!(n.timestamp, timestamp)
+                }
+                NetworkNotification::ConnnectionClosed(n) => {
+                    assert_eq!(PeerId::from_str(&n.peer_id).unwrap(), peer_id);
+                    assert_eq!(Multiaddr::from_str(&n.address).unwrap(), address);
+                    assert_eq!(n.timestamp, timestamp)
+                }
             }
-            _ => panic!("Parsed notification did not matched expected variant"),
         }
     }
 }