From c94e19af56a7875862d43e5443f10332e2fa679b Mon Sep 17 00:00:00 2001
From: immanelg <119798691+immanelg@users.noreply.github.com>
Date: Sat, 6 Jan 2024 18:59:57 +0300
Subject: [PATCH] add test for get + use read() from AsyncReadExt

---
 memcrab-protocol/src/lib.rs | 36 +++++++++++++++++++++++++++++-------
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/memcrab-protocol/src/lib.rs b/memcrab-protocol/src/lib.rs
index ed3d482..c987a37 100644
--- a/memcrab-protocol/src/lib.rs
+++ b/memcrab-protocol/src/lib.rs
@@ -1,5 +1,6 @@
 use std::io;
 use tokio::net::tcp::OwnedReadHalf;
+use tokio::io::AsyncReadExt;
 
 type Bytes = Vec<u8>;
 
@@ -22,8 +23,7 @@ pub trait AsyncReader {
 #[async_trait::async_trait]
 impl AsyncReader for OwnedReadHalf {
     async fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
-        self.readable().await?;
-        self.try_read(buf)
+        AsyncReadExt::read(self, buf).await
     }
 }
 
@@ -120,28 +120,50 @@ mod tests {
     use super::*;
 
     struct MockReader {
-        data: Vec<Vec<u8>>,
+        parts: Vec<Vec<u8>>,
     }
 
     #[async_trait::async_trait]
     impl AsyncReader for MockReader {
         async fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
-            if !self.data.is_empty() {
-                let part = self.data.remove(0);
+            if !self.parts.is_empty() {
+                let part = self.parts.remove(0);
                 for (i, b) in part.iter().enumerate() {
                     buf[i] = *b;
                 }
-                Ok(100)
+                Ok(part.len())
             } else {
                 Ok(0)
             }
         }
     }
 
+    #[tokio::test]
+    async fn test_get() {
+        let mock_reader = MockReader {
+            parts: vec![vec![0, 2, 1, 1].into()],
+        };
+        let mut producer = Producer::new(mock_reader);
+
+        let msg = producer.next_msg().await;
+        assert_eq!(msg.unwrap().unwrap(), Msg::Get(vec![1, 1]))
+    }
+
+    #[tokio::test]
+    async fn test_get_with_partitions() {
+        let mock_reader = MockReader {
+            parts: vec![vec![0].into(), vec![3, 1].into(), vec![2, 3].into()],
+        };
+        let mut producer = Producer::new(mock_reader);
+
+        let msg = producer.next_msg().await;
+        assert_eq!(msg.unwrap().unwrap(), Msg::Get(vec![1, 2, 3]))
+    }
+
     #[tokio::test]
     async fn test_set() {
         let mock_reader = MockReader {
-            data: vec![vec![1, 1, 1, 0, 0, 0, 3, 8, 8, 8]].into(),
+            parts: vec![vec![1, 1, 1, 0, 0, 0, 3, 8, 8, 8]].into(),
         };
         let mut producer = Producer::new(mock_reader);