From 43008b7b073e8ef965fcca908f2b796f72f1b03d Mon Sep 17 00:00:00 2001
From: Sanket Shanbhag <TechieBoy@users.noreply.github.com>
Date: Sat, 24 Dec 2022 14:55:09 +0530
Subject: [PATCH] feat: Add PeerEvent to NetworkEvent enum (#574)

* Add peerAdded and peerRemoved events to NetworkEvent enum

* Remove unnecessary todo macros
---
 crates/net/network/src/manager.rs      | 5 ++++-
 crates/net/network/src/transactions.rs | 2 ++
 crates/net/network/tests/it/connect.rs | 6 ++++++
 crates/net/network/tests/it/testnet.rs | 4 ++--
 4 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/crates/net/network/src/manager.rs b/crates/net/network/src/manager.rs
index 932c1063f4..680cc85055 100644
--- a/crates/net/network/src/manager.rs
+++ b/crates/net/network/src/manager.rs
@@ -51,7 +51,6 @@ use std::{
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tracing::{error, info, trace, warn};
-
 /// Manages the _entire_ state of the network.
 ///
 /// This is an endless [`Future`] that consistently drives the state of the entire network forward.
@@ -676,6 +675,10 @@ pub enum NetworkEvent {
         /// The status of the peer to which a session was established.
         status: Status,
     },
+    /// Event emitted when a new peer is added
+    PeerAdded(PeerId),
+    /// Event emitted when a new peer is removed
+    PeerRemoved(PeerId),
 }
 
 /// Bundles all listeners for [`NetworkEvent`]s.
diff --git a/crates/net/network/src/transactions.rs b/crates/net/network/src/transactions.rs
index 62c6711f2d..d8b0463531 100644
--- a/crates/net/network/src/transactions.rs
+++ b/crates/net/network/src/transactions.rs
@@ -312,6 +312,8 @@ where
                     msg,
                 })
             }
+            // TODO Add remaining events
+            _ => {}
         }
     }
 
diff --git a/crates/net/network/tests/it/connect.rs b/crates/net/network/tests/it/connect.rs
index b568407105..80d908a4f8 100644
--- a/crates/net/network/tests/it/connect.rs
+++ b/crates/net/network/tests/it/connect.rs
@@ -65,6 +65,12 @@ async fn test_establish_connections() {
                 NetworkEvent::SessionEstablished { peer_id, .. } => {
                     assert!(expected_connections.remove(&peer_id))
                 }
+                NetworkEvent::PeerAdded(peer_id) => {
+                    assert!(!expected_connections.contains(&peer_id))
+                }
+                NetworkEvent::PeerRemoved(_) => {
+                    panic!("unexpected event")
+                }
             }
         }
         assert!(expected_connections.is_empty());
diff --git a/crates/net/network/tests/it/testnet.rs b/crates/net/network/tests/it/testnet.rs
index ba0807fdac..4c19200fae 100644
--- a/crates/net/network/tests/it/testnet.rs
+++ b/crates/net/network/tests/it/testnet.rs
@@ -308,7 +308,7 @@ impl NetworkEventStream {
         while let Some(ev) = self.inner.next().await {
             match ev {
                 NetworkEvent::SessionClosed { peer_id, reason } => return Some((peer_id, reason)),
-                NetworkEvent::SessionEstablished { .. } => continue,
+                _ => continue,
             }
         }
         None
@@ -317,8 +317,8 @@ impl NetworkEventStream {
     pub async fn next_session_established(&mut self) -> Option<PeerId> {
         while let Some(ev) = self.inner.next().await {
             match ev {
-                NetworkEvent::SessionClosed { .. } => continue,
                 NetworkEvent::SessionEstablished { peer_id, .. } => return Some(peer_id),
+                _ => continue,
             }
         }
         None