From 9e126fe2cd67e970fe192804f00e51b2caa16edf Mon Sep 17 00:00:00 2001
From: ennucore <ennucore@gmail.com>
Date: Sat, 20 Nov 2021 23:59:01 +0300
Subject: [PATCH] Write receiving in transport (with tests)

---
 src/transport.rs | 41 ++++++++++++++++++++++++++++++++++++++---
 1 file changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/transport.rs b/src/transport.rs
index 6fe63fa..173de39 100644
--- a/src/transport.rs
+++ b/src/transport.rs
@@ -1,5 +1,5 @@
 use alloc::boxed::Box;
-use alloc::string::String;
+use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 use alloc::vec;
 use crate::interface::{Interface, TargetingData};
@@ -98,12 +98,35 @@ impl Transport {
             }
         }
     }
+
+    /// Poll all the interfaces and receive a message
+    ///
+    /// Returns a result with an option of `(message, peer_id)`
+    pub fn receive(&mut self) -> Option<(MessageBytes, u64 /* peer id*/)> {
+        if let Some((interface_id, (msg, peer_data))) = self.interfaces
+            .iter_mut()
+            // For each interface return (interface id, message result)
+            .map(|interface| (interface.id().to_string(), interface.receive()))
+            // If there was an error, print it
+            .map(|res| match res {
+                (id, Err(e)) => {
+                    #[cfg(std)]
+                    println!("An error occurred while receiving: {:?}", e);
+                    (id, Err(e))
+                }
+                (id, Ok(r)) => (id, Ok(r))
+            })
+            // Find a result where there is a message
+            .find(|r| matches!(r, (_, Ok(Some(_)))))
+            // Safely unwrap this result (we already matched `Ok(Some(_))`)
+            .map(|(id, r)| (id, r.unwrap().unwrap())) {
+            Some((msg, self.find_or_add_peer(interface_id, peer_data)))
+        } else { None }
+    }
 }
 
 #[cfg(test)]
 use crate::interface::test_interface::TestInterface;
-#[cfg(test)]
-use alloc::string::ToString;
 
 #[test]
 fn test_adding_peer_to_transport() {
@@ -131,3 +154,15 @@ fn test_transport_sending() {
     transport.send_message(vec![239, 123], Some(peer_id)).unwrap();
     assert_eq!(transport.interfaces[0].receive(), IFResult::Ok(Some((vec![239, 123], interface_targeting_data))));
 }
+
+#[test]
+fn test_transport_receiving() {
+    let mut transport = Transport::new(vec![Box::new(TestInterface::default())]);
+    let (interface_id, interface_targeting_data) = ("test_interface".to_string(), "hi".to_string());
+    let peer_id = transport.find_or_add_peer(interface_id.clone(), interface_targeting_data.clone());
+    transport.send_message(vec![239, 123], None).unwrap();
+    assert_eq!(transport.receive(), Some((vec![239u8, 123], transport.get_peer_by_parameters(interface_id.as_str(), "").unwrap().peer_id)));
+    assert!(transport.receive().is_none());
+    transport.send_message(vec![239, 123], Some(peer_id)).unwrap();
+    assert_eq!(transport.receive(), Some((vec![239, 123], peer_id)));
+}