[internal] test batch and stream API paths for load/store from remote store

src/rust/engine/fs/store/src/remote.rs

@@ -235,7 +235,7 @@ impl ByteStore {
     }
   }
 
-  async fn store_bytes_source_batch<ByteSource>(
+  pub(crate) async fn store_bytes_source_batch<ByteSource>(
     &self,
     digest: Digest,
     bytes: ByteSource,
@@ -259,7 +259,7 @@ impl ByteStore {
     Ok(())
   }
 
-  async fn store_bytes_source_stream<ByteSource>(
+  pub(crate) async fn store_bytes_source_stream<ByteSource>(
     &self,
     digest: Digest,
     bytes: ByteSource,
@@ -349,7 +349,7 @@ impl ByteStore {
     }
   }
 
-  async fn load_bytes_with_batch<
+  pub(crate) async fn load_bytes_with_batch<
     T: Send + 'static,
     F: Fn(Bytes) -> Result<T, String> + Send + Sync + Clone + 'static,
   >(
@@ -390,7 +390,7 @@ impl ByteStore {
     }
   }
 
-  async fn load_bytes_with_stream<
+  pub(crate) async fn load_bytes_with_stream<
     T: Send + 'static,
     F: Fn(Bytes) -> Result<T, String> + Send + Sync + Clone + 'static,
   >(

src/rust/engine/fs/store/src/remote_tests.rs

@@ -307,6 +307,63 @@ async fn list_missing_digests_error() {
   );
 }
 
+#[tokio::test]
+async fn load_via_each_api() {
+  let _ = WorkunitStore::setup_for_tests();
+  let cas = StubCAS::empty();
+
+  let data = TestData::roland();
+  {
+    let mut blobs = cas.blobs.lock();
+    blobs.insert(data.fingerprint(), data.bytes());
+  }
+
+  let store = new_byte_store(&cas);
+  let result_batch = store
+    .load_bytes_with_batch(data.digest(), |b| Ok(b))
+    .await
+    .unwrap()
+    .unwrap();
+  let result_stream = store
+    .load_bytes_with_stream(data.digest(), |b| Ok(b))
+    .await
+    .unwrap()
+    .unwrap();
+  assert_eq!(result_batch, data.bytes());
+  assert_eq!(result_stream, data.bytes());
+}
+
+#[tokio::test]
+async fn store_via_each_api() {
+  let _ = WorkunitStore::setup_for_tests();
+  let cas = StubCAS::empty();
+
+  let data = TestData::roland();
+  let store = new_byte_store(&cas);
+
+  let bytes = data.bytes();
+  let _ = store
+    .store_bytes_source_batch(data.digest(), move |r| bytes.slice(r))
+    .await
+    .unwrap();
+  {
+    let mut blobs = cas.blobs.lock();
+    assert_eq!(*blobs.get(&data.digest().hash).unwrap(), data.bytes());
+    blobs.clear();
+  }
+
+  let bytes = data.bytes();
+  let _ = store
+    .store_bytes_source_stream(data.digest(), move |r| bytes.slice(r))
+    .await
+    .unwrap();
+  {
+    let mut blobs = cas.blobs.lock();
+    assert_eq!(*blobs.get(&data.digest().hash).unwrap(), data.bytes());
+    blobs.clear();
+  }
+}
+
 fn new_byte_store(cas: &StubCAS) -> ByteStore {
   ByteStore::new(
     &cas.address(),