adding a mocked database test

akd/src/tests.rs

@@ -22,7 +22,8 @@ use crate::{
         types::{DbRecord, KeyData, ValueState, ValueStateRetrievalFlag},
         Database, DbSetState, Storable,
     },
-    AkdLabel, AkdValue, HistoryParams, HistoryVerificationParams, VerifyResult,
+    tree_node::TreeNodeWithPreviousValue,
+    AkdLabel, AkdValue, Azks, HistoryParams, HistoryVerificationParams, VerifyResult,
 };
 
 #[derive(Clone)]
@@ -80,6 +81,76 @@ mockall::mock! {
     }
 }
 
+fn setup_mocked_db(db: &mut MockLocalDatabase, test_db: &AsyncInMemoryDatabase) {
+    let tokio_handle = tokio::runtime::Handle::current();
+
+    // ===== Set ===== //
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_set()
+        .returning(move |record| handle.block_on(tmp_db.set(record)));
+
+    // ===== Batch Set ===== //
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_batch_set()
+        .returning(move |record, other| handle.block_on(tmp_db.batch_set(record, other)));
+
+    // ===== Get ===== //
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_get::<Azks>()
+        .returning(move |key| handle.block_on(tmp_db.get::<Azks>(key)));
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_get::<TreeNodeWithPreviousValue>()
+        .returning(move |key| handle.block_on(tmp_db.get::<TreeNodeWithPreviousValue>(key)));
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_get::<Azks>()
+        .returning(move |key| handle.block_on(tmp_db.get::<Azks>(key)));
+
+    // ===== Batch Get ===== //
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_batch_get::<Azks>()
+        .returning(move |key| handle.block_on(tmp_db.batch_get::<Azks>(key)));
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_batch_get::<TreeNodeWithPreviousValue>()
+        .returning(move |key| handle.block_on(tmp_db.batch_get::<TreeNodeWithPreviousValue>(key)));
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_batch_get::<Azks>()
+        .returning(move |key| handle.block_on(tmp_db.batch_get::<Azks>(key)));
+
+    // ===== Get User Data ===== //
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_get_user_data()
+        .returning(move |arg| handle.block_on(tmp_db.get_user_data(arg)));
+
+    // ===== Get User State ===== //
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_get_user_state()
+        .returning(move |arg, flag| handle.block_on(tmp_db.get_user_state(arg, flag)));
+
+    // ===== Get User State Versions ===== //
+    let handle = tokio_handle.clone();
+    let tmp_db = test_db.clone();
+    db.expect_get_user_state_versions()
+        .returning(move |arg, flag| handle.block_on(tmp_db.get_user_state_versions(arg, flag)));
+
+
+    // // ===== Clone ===== //
+    // db.expect_clone().return_const(|| {
+    //     let db = 
+    // });
+
+
+}
+
 // 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> {
@@ -1119,6 +1190,51 @@ async fn test_tombstoned_key_history() -> Result<(), AkdError> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_publish_op_makes_no_get_requests() {
+    let test_db = AsyncInMemoryDatabase::new();
+
+    let mut db = MockLocalDatabase {
+        ..Default::default()
+    };
+    setup_mocked_db(&mut db, &test_db);
+
+    let storage = StorageManager::new_no_cache(db.clone());
+    let vrf = HardCodedAkdVRF {};
+    let akd = Directory::<_, _>::new(storage, vrf, false)
+        .await
+        .expect("Failed to create directory");
+
+    // Create a set with 2 updates, (label, value) pairs
+    // ("hello10", "hello10")
+    // ("hello11", "hello11")
+    let mut updates = vec![];
+    for i in 0..1 {
+        updates.push((
+            AkdLabel(format!("hello1{}", i).as_bytes().to_vec()),
+            AkdValue(format!("hello1{}", i).as_bytes().to_vec()),
+        ));
+    }
+    // Publish the updates. Now the akd's epoch will be 1.
+    akd.publish(updates).await.expect("Failed to do initial publish");
+
+    db.expect_get::<TreeNodeWithPreviousValue>().returning(|_| Err(StorageError::Other("Boom!".to_string())));
+
+    // create more updates
+    let mut updates = vec![];
+    for i in 0..1 {
+        updates.push((
+            AkdLabel(format!("hello1{}", i).as_bytes().to_vec()),
+            AkdValue(format!("hello1{}", i+1).as_bytes().to_vec()),
+        ));
+    }
+
+    // try to publish again, this time with the "boom" returning from any mocked get-calls
+    // on tree nodes
+    akd.publish(updates).await.expect("Failed to do subsequent publish");
+
+}
+
 // Test coverage on issue #144, verification failures with
 // small trees (<4 nodes) in both the tests below
 // Note that the use of a VRF means that that the label