From f990b2d1a5445ebd67ef51be33c17dc270726321 Mon Sep 17 00:00:00 2001
From: Alexander Schmidt <alexander.schmidt@iota.org>
Date: Tue, 19 Oct 2021 13:19:55 +0200
Subject: [PATCH] Add apply_condition unit test

---
 bee-autopeering/src/filter.rs      | 21 ++++++++++++++++++---
 bee-autopeering/src/lib.rs         |  4 ++--
 bee-autopeering/src/peer.rs        | 30 ++++++++++++++++++++++++++++--
 bee-autopeering/src/service_map.rs | 21 +++++++++++++++++++++
 4 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/bee-autopeering/src/filter.rs b/bee-autopeering/src/filter.rs
index 42c3aec061..a37d004335 100644
--- a/bee-autopeering/src/filter.rs
+++ b/bee-autopeering/src/filter.rs
@@ -64,7 +64,7 @@ impl Filter {
 
 #[cfg(test)]
 mod tests {
-    use crate::service_map::AUTOPEERING_SERVICE_NAME;
+    use crate::service_map::{ServiceProtocol, AUTOPEERING_SERVICE_NAME};
 
     use super::*;
 
@@ -75,18 +75,33 @@ mod tests {
     }
 
     #[test]
-    fn add_condition() {
+    fn apply_condition() {
         let mut filter = Filter::new();
         assert_eq!(0, filter.num_conditions());
 
-        let condition = |peer: &Peer| -> bool { peer.services().port(AUTOPEERING_SERVICE_NAME).unwrap() == 8080 };
+        let condition = |peer: &Peer| -> bool { peer.services().port(AUTOPEERING_SERVICE_NAME).unwrap() == 1337 };
         let condition = Box::new(condition);
 
         filter.add_condition(condition);
         assert_eq!(1, filter.num_conditions());
 
+        let mut peer1 = Peer::new_test_peer(1);
+        peer1.add_service(AUTOPEERING_SERVICE_NAME, ServiceProtocol::Udp, 6969);
+        assert_eq!(1, peer1.num_services());
+
+        let mut peer2 = Peer::new_test_peer(2);
+        peer2.add_service(AUTOPEERING_SERVICE_NAME, ServiceProtocol::Udp, 1337);
+        assert_eq!(1, peer2.num_services());
+
+        let candidates = [peer1, peer2];
+        let filtered = filter.apply(&candidates);
+        assert_eq!(1, filtered.len());
+
         filter.clear_conditions();
         assert_eq!(0, filter.num_conditions());
+
+        let filtered = filter.apply(&candidates);
+        assert_eq!(2, filtered.len());
     }
 
     #[test]
diff --git a/bee-autopeering/src/lib.rs b/bee-autopeering/src/lib.rs
index 5a564468af..f968d15df0 100644
--- a/bee-autopeering/src/lib.rs
+++ b/bee-autopeering/src/lib.rs
@@ -24,16 +24,16 @@ mod proto {
 mod request;
 mod salt;
 mod server;
-mod service_map;
 mod time;
 
 pub mod identity;
 pub mod init;
 pub mod peer;
 pub mod peerstore;
+pub mod service_map;
 
 pub use identity::PeerId;
 pub use init::init;
 pub use local::Local;
 pub use peer::Peer;
-pub use service_map::{ServiceMap, ServiceName};
+pub use service_map::{ServiceMap, ServiceName, ServiceProtocol};
diff --git a/bee-autopeering/src/peer.rs b/bee-autopeering/src/peer.rs
index c9f7cbb6c9..15768772a8 100644
--- a/bee-autopeering/src/peer.rs
+++ b/bee-autopeering/src/peer.rs
@@ -1,14 +1,22 @@
 // Copyright 2021 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{identity::PeerId, proto, service_map::ServiceMap};
+use crate::{
+    identity::PeerId,
+    proto,
+    service_map::{ServiceMap, ServiceProtocol},
+};
 
 use bytes::BytesMut;
 use crypto::signatures::ed25519::PublicKey;
 use prost::{DecodeError, EncodeError, Message};
 // use serde::{Deserialize, Serialize};
 
-use std::{convert::TryInto, fmt, net::IpAddr};
+use std::{
+    convert::TryInto,
+    fmt,
+    net::{IpAddr, Ipv4Addr},
+};
 
 /// Represents a peer.
 // #[derive(Serialize, Deserialize)]
@@ -49,6 +57,20 @@ impl Peer {
         &self.services
     }
 
+    pub fn add_service(&mut self, service_name: impl ToString, protocol: ServiceProtocol, port: u16) {
+        let ipv = match self.ip_address {
+            IpAddr::V4(ip4) => "ip4",
+            IpAddr::V6(ip6) => "ip6",
+        };
+
+        // Unwrap: we control what is being parsed.
+        let multiaddr = format!("/{}/{}/{}/{}", ipv, self.ip_address, protocol, port)
+            .parse()
+            .unwrap();
+
+        self.services.insert(service_name.to_string(), multiaddr);
+    }
+
     /// Creates a discovered peer from its Protobuf representation/encoding.
     pub fn from_protobuf(bytes: &[u8]) -> Result<Self, DecodeError> {
         Ok(proto::Peer::decode(bytes)?.into())
@@ -143,5 +165,9 @@ mod tests {
                 services,
             }
         }
+
+        pub(crate) fn num_services(&self) -> usize {
+            self.services().len()
+        }
     }
 }
diff --git a/bee-autopeering/src/service_map.rs b/bee-autopeering/src/service_map.rs
index b61a07ece1..e189d83be7 100644
--- a/bee-autopeering/src/service_map.rs
+++ b/bee-autopeering/src/service_map.rs
@@ -133,11 +133,32 @@ impl fmt::Display for ServiceMap {
     }
 }
 
+pub enum ServiceProtocol {
+    Tcp,
+    Udp,
+}
+
+impl fmt::Display for ServiceProtocol {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let protocol = match self {
+            ServiceProtocol::Udp => "udp",
+            ServiceProtocol::Tcp => "tcp",
+        };
+        write!(f, "{}", protocol)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::proto;
 
+    impl ServiceMap {
+        pub(crate) fn len(&self) -> usize {
+            self.0.len()
+        }
+    }
+
     #[test]
     fn convert_service_map() {
         let mut map = HashMap::new();