Adding mock support for the Database trait in AKD

Related to: facebook#332
slawlor · Jan 3, 2023 · 1e0d5d9 · 1e0d5d9
1 parent 135ffbe
commit 1e0d5d9
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 4 deletions.
diff --git a/akd/Cargo.toml b/akd/Cargo.toml
@@ -62,6 +62,7 @@ once_cell = { version = "1" }
 ctor = "0.1"
 tokio-test = "0.4"
 tokio = { version = "1.21", features = ["rt", "sync", "time", "macros"] }
+mockall = "0.11"
 
 # To enable the public-test feature in tests
 akd = { path = ".", features = ["public-tests"], default-features = false }

diff --git a/akd/src/storage/mod.rs b/akd/src/storage/mod.rs
@@ -43,7 +43,7 @@ pub enum DbSetState {
 
 /// Storable represents an _item_ which can be stored in the storage layer
 #[cfg(feature = "serde_serialization")]
-pub trait Storable: Clone + Serialize + DeserializeOwned + Sync {
+pub trait Storable: Clone + Serialize + DeserializeOwned + Sync + 'static {
     /// This particular storage will have a key type
     type StorageKey: Clone + Serialize + Eq + Hash + Send + Sync + std::fmt::Debug;
 
@@ -68,7 +68,7 @@ pub trait Storable: Clone + Serialize + DeserializeOwned + Sync {
 
 /// Storable represents an _item_ which can be stored in the storage layer
 #[cfg(not(feature = "serde_serialization"))]
-pub trait Storable: Clone + Sync {
+pub trait Storable: Clone + Sync + 'static {
     /// This particular storage will have a key type
     type StorageKey: Clone + Eq + Hash + Send + Sync + std::fmt::Debug;
 

diff --git a/akd/src/tests.rs b/akd/src/tests.rs
@@ -8,16 +8,78 @@
 
 //! Contains the tests for the high-level API (directory, auditor, client)
 
+use std::collections::HashMap;
+
 use crate::{
     auditor::audit_verify,
     client::{key_history_verify, lookup_verify},
     directory::{Directory, PublishCorruption},
     ecvrf::{HardCodedAkdVRF, VRFKeyStorage},
-    errors::AkdError,
-    storage::{manager::StorageManager, memory::AsyncInMemoryDatabase, types::DbRecord, Database},
+    errors::{AkdError, StorageError},
+    storage::{
+        manager::StorageManager,
+        memory::AsyncInMemoryDatabase,
+        types::{DbRecord, KeyData, ValueState, ValueStateRetrievalFlag},
+        Database, DbSetState, Storable,
+    },
     AkdLabel, AkdValue, HistoryParams, HistoryVerificationParams, VerifyResult,
 };
 
+#[derive(Clone)]
+pub struct LocalDatabase;
+
+unsafe impl Send for LocalDatabase {}
+unsafe impl Sync for LocalDatabase {}
+
+mockall::mock! {
+    pub LocalDatabase {
+
+    }
+    impl Clone for LocalDatabase {
+        fn clone(&self) -> Self;
+    }
+    #[async_trait::async_trait]
+    impl Database for LocalDatabase {
+        /// Set a record in the database
+        async fn set(&self, record: DbRecord) -> Result<(), StorageError>;
+
+        /// Set multiple records in the database with a minimal set of operations
+        async fn batch_set(
+            &self,
+            records: Vec<DbRecord>,
+            state: DbSetState,
+        ) -> Result<(), StorageError>;
+
+        /// Retrieve a stored record from the database
+        async fn get<St: Storable>(&self, id: &St::StorageKey) -> Result<DbRecord, StorageError>;
+
+        /// Retrieve a batch of records by id from the database
+        async fn batch_get<St: Storable>(
+            &self,
+            ids: &[St::StorageKey],
+        ) -> Result<Vec<DbRecord>, StorageError>;
+
+        /* User data searching */
+
+        /// Retrieve the user data for a given user
+        async fn get_user_data(&self, username: &AkdLabel) -> Result<KeyData, StorageError>;
+
+        /// Retrieve a specific state for a given user
+        async fn get_user_state(
+            &self,
+            username: &AkdLabel,
+            flag: ValueStateRetrievalFlag,
+        ) -> Result<ValueState, StorageError>;
+
+        /// Retrieve the user -> state version mapping in bulk. This is the same as get_user_states but with less data retrieved from the storage layer
+        async fn get_user_state_versions(
+            &self,
+            usernames: &[AkdLabel],
+            flag: ValueStateRetrievalFlag,
+        ) -> Result<HashMap<AkdLabel, (u64, AkdValue)>, StorageError>;
+    }
+}
+
 // A simple test to ensure that the empty tree hashes to the correct value
 #[tokio::test]
 async fn test_empty_tree_root_hash() -> Result<(), AkdError> {