From 51546eb3875663a3f38fe015a727cda4fbdee43c Mon Sep 17 00:00:00 2001
From: Edward Shen <code@eddie.sh>
Date: Fri, 16 Jul 2021 14:06:28 -0400
Subject: [PATCH] Add memory cache get tests

---
 src/cache/mem.rs | 104 ++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 93 insertions(+), 11 deletions(-)

diff --git a/src/cache/mem.rs b/src/cache/mem.rs
index 82bd3fc..0f8e6ee 100644
--- a/src/cache/mem.rs
+++ b/src/cache/mem.rs
@@ -202,6 +202,7 @@ where
 
 #[cfg(test)]
 mod test_util {
+    use std::cell::RefCell;
     use std::collections::HashMap;
 
     use super::{CacheValue, InternalMemoryCache};
@@ -209,10 +210,14 @@ mod test_util {
         Cache, CacheEntry, CacheError, CacheKey, CacheStream, CallbackCache, ImageMetadata,
     };
     use async_trait::async_trait;
+    use parking_lot::Mutex;
+    use tokio::io::BufReader;
     use tokio::sync::mpsc::Sender;
+    use tokio_util::codec::{BytesCodec, FramedRead};
 
+    #[derive(Default)]
     pub struct TestDiskCache(
-        pub HashMap<CacheKey, Result<(CacheStream, ImageMetadata), CacheError>>,
+        pub Mutex<RefCell<HashMap<CacheKey, Result<(CacheStream, ImageMetadata), CacheError>>>>,
     );
 
     #[async_trait]
@@ -221,8 +226,7 @@ mod test_util {
             &self,
             key: &CacheKey,
         ) -> Option<Result<(CacheStream, ImageMetadata), CacheError>> {
-            // todo: Actually implement nontrivial code
-            None
+            self.0.lock().get_mut().remove(key)
         }
 
         async fn put(
@@ -231,7 +235,12 @@ mod test_util {
             image: bytes::Bytes,
             metadata: ImageMetadata,
         ) -> Result<(), CacheError> {
-            todo!()
+            let reader = Box::pin(BufReader::new(tokio_util::io::StreamReader::new(
+                tokio_stream::once(Ok::<_, std::io::Error>(image)),
+            )));
+            let stream = CacheStream::Completed(FramedRead::new(reader, BytesCodec::new()));
+            self.0.lock().get_mut().insert(key, Ok((stream, metadata)));
+            Ok(())
         }
     }
 
@@ -272,11 +281,10 @@ mod test_util {
 
 #[cfg(test)]
 mod cache_ops {
-    use std::collections::HashMap;
     use std::error::Error;
 
     use bytes::Bytes;
-    use futures::FutureExt;
+    use futures::{FutureExt, StreamExt};
 
     use crate::cache::mem::InternalMemoryCache;
     use crate::cache::{Cache, CacheKey, CacheStream, ImageMetadata, MemStream};
@@ -287,7 +295,7 @@ mod cache_ops {
     #[tokio::test]
     async fn get_mem_cached() -> Result<(), Box<dyn Error>> {
         let (cache, mut rx) = MemoryCache::<TestMemoryCache, _>::new_with_receiver(
-            TestDiskCache(HashMap::new()),
+            TestDiskCache::default(),
             crate::units::Bytes(10),
         );
 
@@ -320,11 +328,85 @@ mod cache_ops {
         Ok(())
     }
 
-    #[test]
-    fn get_disk_cached() {}
+    #[tokio::test]
+    async fn get_disk_cached() -> Result<(), Box<dyn Error>> {
+        let (mut cache, mut rx) = MemoryCache::<TestMemoryCache, _>::new_with_receiver(
+            TestDiskCache::default(),
+            crate::units::Bytes(10),
+        );
 
-    #[test]
-    fn get_miss() {}
+        let key = CacheKey("a".to_string(), "b".to_string(), false);
+        let metadata = ImageMetadata {
+            content_type: None,
+            content_length: Some(1),
+            last_modified: None,
+        };
+        let bytes = Bytes::from_static(b"abcd");
+
+        {
+            let cache = &mut cache.inner;
+            cache
+                .put(key.clone(), bytes.clone(), metadata.clone())
+                .await?;
+        }
+
+        let (mut stream, ret_metadata) = cache.get(&key).await.unwrap()?;
+        assert_eq!(metadata, ret_metadata);
+        assert!(matches!(stream, CacheStream::Completed(_)));
+        assert_eq!(stream.next().await, Some(Ok(bytes.clone())));
+
+        assert!(rx.recv().now_or_never().is_none());
+
+        Ok(())
+    }
+
+    // Identical to the get_disk_cached test but we hold a lock on the mem_cache
+    #[tokio::test]
+    async fn get_mem_locked() -> Result<(), Box<dyn Error>> {
+        let (mut cache, mut rx) = MemoryCache::<TestMemoryCache, _>::new_with_receiver(
+            TestDiskCache::default(),
+            crate::units::Bytes(10),
+        );
+
+        let key = CacheKey("a".to_string(), "b".to_string(), false);
+        let metadata = ImageMetadata {
+            content_type: None,
+            content_length: Some(1),
+            last_modified: None,
+        };
+        let bytes = Bytes::from_static(b"abcd");
+
+        {
+            let cache = &mut cache.inner;
+            cache
+                .put(key.clone(), bytes.clone(), metadata.clone())
+                .await?;
+        }
+
+        // intentionally not dropped
+        let _mem_cache = &mut cache.mem_cache.lock().await;
+
+        let (mut stream, ret_metadata) = cache.get(&key).await.unwrap()?;
+        assert_eq!(metadata, ret_metadata);
+        assert!(matches!(stream, CacheStream::Completed(_)));
+        assert_eq!(stream.next().await, Some(Ok(bytes.clone())));
+
+        assert!(rx.recv().now_or_never().is_none());
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn get_miss() {
+        let (cache, mut rx) = MemoryCache::<TestMemoryCache, _>::new_with_receiver(
+            TestDiskCache::default(),
+            crate::units::Bytes(10),
+        );
+
+        let key = CacheKey("a".to_string(), "b".to_string(), false);
+        assert!(cache.get(&key).await.is_none());
+        assert!(rx.recv().now_or_never().is_none());
+    }
 }
 
 #[cfg(test)]