1 files changed, 20 insertions, 26 deletions

diff --git a/tvix/castore/src/proto/tests/grpc_blobservice.rs b/tvix/castore/src/proto/tests/grpc_blobservice.rs
index 0d7b340b4409..a7816bd1e961 100644
--- a/tvix/castore/src/proto/tests/grpc_blobservice.rs
+++ b/tvix/castore/src/proto/tests/grpc_blobservice.rs
@@ -1,23 +1,17 @@
 use crate::fixtures::{BLOB_A, BLOB_A_DIGEST};
-use crate::proto::blob_service_server::BlobService as GRPCBlobService;
-use crate::proto::{BlobChunk, GRPCBlobServiceWrapper, ReadBlobRequest, StatBlobRequest};
-use crate::utils::gen_blob_service;
+use crate::proto::{BlobChunk, ReadBlobRequest, StatBlobRequest};
+use crate::utils::gen_blobsvc_grpc_client;
 use tokio_stream::StreamExt;
 
-fn gen_grpc_blob_service() -> GRPCBlobServiceWrapper {
-    let blob_service = gen_blob_service();
-    GRPCBlobServiceWrapper::from(blob_service)
-}
-
 /// Trying to read a non-existent blob should return a not found error.
 #[tokio::test]
 async fn not_found_read() {
-    let service = gen_grpc_blob_service();
+    let mut grpc_client = gen_blobsvc_grpc_client().await;
 
-    let resp = service
-        .read(tonic::Request::new(ReadBlobRequest {
+    let resp = grpc_client
+        .read(ReadBlobRequest {
             digest: BLOB_A_DIGEST.clone().into(),
-        }))
+        })
         .await;
 
     // We can't use unwrap_err here, because the Ok value doesn't implement
@@ -32,13 +26,13 @@ async fn not_found_read() {
 /// Trying to stat a non-existent blob should return a not found error.
 #[tokio::test]
 async fn not_found_stat() {
-    let service = gen_grpc_blob_service();
+    let mut grpc_client = gen_blobsvc_grpc_client().await;
 
-    let resp = service
-        .stat(tonic::Request::new(StatBlobRequest {
+    let resp = grpc_client
+        .stat(StatBlobRequest {
             digest: BLOB_A_DIGEST.clone().into(),
             ..Default::default()
-        }))
+        })
         .await
         .expect_err("must fail");
 
@@ -49,13 +43,13 @@ async fn not_found_stat() {
 /// Put a blob in the store, get it back.
 #[tokio::test]
 async fn put_read_stat() {
-    let service = gen_grpc_blob_service();
+    let mut grpc_client = gen_blobsvc_grpc_client().await;
 
     // Send blob A.
-    let put_resp = service
-        .put(tonic_mock::streaming_request(vec![BlobChunk {
+    let put_resp = grpc_client
+        .put(tokio_stream::once(BlobChunk {
             data: BLOB_A.clone(),
-        }]))
+        }))
         .await
         .expect("must succeed")
         .into_inner();
@@ -65,20 +59,20 @@ async fn put_read_stat() {
     // Stat for the digest of A.
     // We currently don't ask for more granular chunking data, as we don't
     // expose it yet.
-    let _resp = service
-        .stat(tonic::Request::new(StatBlobRequest {
+    let _resp = grpc_client
+        .stat(StatBlobRequest {
             digest: BLOB_A_DIGEST.clone().into(),
             ..Default::default()
-        }))
+        })
         .await
         .expect("must succeed")
         .into_inner();
 
     // Read the blob. It should return the same data.
-    let resp = service
-        .read(tonic::Request::new(ReadBlobRequest {
+    let resp = grpc_client
+        .read(ReadBlobRequest {
             digest: BLOB_A_DIGEST.clone().into(),
-        }))
+        })
         .await;
 
     let mut rx = resp.ok().unwrap().into_inner();