From afbf648528a4c61192bed4e3d033a4a14a97551b Mon Sep 17 00:00:00 2001
From: Dylan Knutson <tcdknutson@gmail.com>
Date: Thu, 25 Apr 2024 00:26:47 -0700
Subject: [PATCH] unit tests for store handler

---
 src/handlers/store_handler.rs | 53 +++++++++++++++++++++++++++++++++++
 src/main.rs                   |  2 +-
 src/sha256.rs                 |  3 +-
 src/shard/mod.rs              |  4 +--
 src/shards.rs                 |  7 +++--
 5 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/src/handlers/store_handler.rs b/src/handlers/store_handler.rs
index 688350e..dc83d3f 100644
--- a/src/handlers/store_handler.rs
+++ b/src/handlers/store_handler.rs
@@ -30,3 +30,56 @@ pub async fn store_handler(
         .store(sha256, request.content_type, request.data.contents)
         .await
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::shard::test::make_shard;
+    use axum::body::Bytes;
+    use axum_typed_multipart::FieldData;
+
+    async fn make_shards() -> Shards {
+        Shards::new(vec![make_shard().await]).unwrap()
+    }
+
+    async fn send_request(sha256: Option<Sha256>, data: Bytes) -> StoreResult {
+        store_handler(
+            Extension(make_shards().await),
+            TypedMultipart(StoreRequest {
+                sha256: sha256.map(|s| s.hex_string()),
+                content_type: "text/plain".to_string(),
+                data: FieldData {
+                    metadata: Default::default(),
+                    contents: data,
+                },
+            }),
+        )
+        .await
+    }
+
+    #[tokio::test]
+    async fn test_store_handler() {
+        let result = send_request(None, "hello, world!".as_bytes().into()).await;
+        assert!(matches!(result, StoreResult::Created { .. }));
+    }
+
+    #[tokio::test]
+    async fn test_store_handler_mismatched_sha256() {
+        let not_hello_world = Sha256::from_bytes("not hello, world!".as_bytes());
+        let hello_world = Sha256::from_bytes("hello, world!".as_bytes());
+        let result = send_request(Some(not_hello_world), "hello, world!".as_bytes().into()).await;
+        assert_eq!(
+            result,
+            StoreResult::Sha256Mismatch {
+                sha256: hello_world.hex_string()
+            }
+        );
+    }
+
+    #[tokio::test]
+    async fn test_store_handler_matching_sha256() {
+        let hello_world = Sha256::from_bytes("hello, world!".as_bytes());
+        let result = send_request(Some(hello_world), "hello, world!".as_bytes().into()).await;
+        assert!(matches!(result, StoreResult::Created { .. }));
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 5c7e14d..1507a8c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -77,7 +77,7 @@ fn main() -> Result<(), Box<dyn Error>> {
             shards_vec.push(shard);
         }
 
-        let shards = Shards::new(shards_vec);
+        let shards = Shards::new(shards_vec).ok_or("num shards must be > 0")?;
         server_loop(server, shards.clone()).await?;
         info!("shutting down server...");
         shards.close_all().await?;
diff --git a/src/sha256.rs b/src/sha256.rs
index add8195..e5d76d8 100644
--- a/src/sha256.rs
+++ b/src/sha256.rs
@@ -42,7 +42,8 @@ impl Sha256 {
         format!("{:x}", self)
     }
     pub fn modulo(&self, num: usize) -> usize {
-        self.0[0] as usize % num
+        let sum: usize = self.0.iter().map(|v| *v as usize).sum();
+        sum % num
     }
 }
 
diff --git a/src/shard/mod.rs b/src/shard/mod.rs
index 2716e43..dcd4949 100644
--- a/src/shard/mod.rs
+++ b/src/shard/mod.rs
@@ -154,10 +154,10 @@ async fn get_num_entries(conn: &Connection) -> Result<usize, tokio_rusqlite::Err
 }
 
 #[cfg(test)]
-mod test {
+pub mod test {
     use crate::sha256::Sha256;
 
-    async fn make_shard() -> super::Shard {
+    pub async fn make_shard() -> super::Shard {
         let conn = tokio_rusqlite::Connection::open_in_memory().await.unwrap();
         super::Shard::open(0, conn).await.unwrap()
     }
diff --git a/src/shards.rs b/src/shards.rs
index 8b54615..1d0ad8e 100644
--- a/src/shards.rs
+++ b/src/shards.rs
@@ -5,8 +5,11 @@ use crate::{sha256::Sha256, shard::Shard};
 #[derive(Clone)]
 pub struct Shards(Vec<Shard>);
 impl Shards {
-    pub fn new(shards: Vec<Shard>) -> Self {
-        Self(shards)
+    pub fn new(shards: Vec<Shard>) -> Option<Self> {
+        if shards.is_empty() {
+            return None;
+        }
+        Some(Self(shards))
     }
 
     pub fn shard_for(&self, sha256: &Sha256) -> &Shard {