diff options
author | Florian Klink <flokli@flokli.de> | 2023-01-18T20·43+0100 |
---|---|---|
committer | flokli <flokli@flokli.de> | 2023-01-21T10·48+0000 |
commit | e719da53be01151a2a232c55b9f74d7fa09224a9 (patch) | |
tree | c817117a2e49caf8febb82374b623736185e968a /tvix/store/src/tests/directory_service.rs | |
parent | 22b9e6ff092c531ee72037ff8f4c065eaff3b839 (diff) |
feat(tvix/store): add SledDirectoryService r/5722
Change-Id: I8e97a59da2c6623f4cca746efa741097a6dbe5b1 Reviewed-on: https://cl.tvl.fyi/c/depot/+/7871 Reviewed-by: tazjin <tazjin@tvl.su> Tested-by: BuildkiteCI
Diffstat (limited to 'tvix/store/src/tests/directory_service.rs')
-rw-r--r-- | tvix/store/src/tests/directory_service.rs | 158 |
1 files changed, 158 insertions, 0 deletions
diff --git a/tvix/store/src/tests/directory_service.rs b/tvix/store/src/tests/directory_service.rs new file mode 100644 index 000000000000..c16feb243066 --- /dev/null +++ b/tvix/store/src/tests/directory_service.rs @@ -0,0 +1,158 @@ +use tempfile::TempDir; +use tokio_stream::StreamExt; +use tonic::Status; + +use crate::proto::directory_service_server::DirectoryService; +use crate::proto::get_directory_request::ByWhat; +use crate::proto::GetDirectoryRequest; +use crate::proto::{Directory, DirectoryNode}; +use crate::sled_directory_service::SledDirectoryService; +use lazy_static::lazy_static; + +lazy_static! { + static ref DIRECTORY_A: Directory = Directory::default(); + static ref DIRECTORY_B: Directory = Directory { + directories: vec![DirectoryNode { + name: "a".to_string(), + digest: DIRECTORY_A.digest(), + size: DIRECTORY_A.size(), + }], + ..Default::default() + }; +} + +/// Send the specified GetDirectoryRequest. +/// Returns an error in the case of an error response, or an error in one of the items in the stream, +/// or a Vec<Directory> in the case of a successful request. +async fn get_directories<S: DirectoryService>( + svc: &S, + get_directory_request: GetDirectoryRequest, +) -> Result<Vec<Directory>, Status> { + let resp = svc.get(tonic::Request::new(get_directory_request)).await; + + // if the response is an error itself, return the error, otherwise unpack + let stream = match resp { + Ok(resp) => resp, + Err(status) => return Err(status), + } + .into_inner(); + + let directory_results: Vec<Result<Directory, Status>> = stream.collect().await; + + // turn Vec<Result<Directory, Status> into Result<Vec<Directory>,Status> + directory_results.into_iter().collect() +} + +/// Trying to get a non-existent Directory should return a not found error. +#[tokio::test] +async fn not_found() -> anyhow::Result<()> { + let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?; + + let resp = service + .get(tonic::Request::new(GetDirectoryRequest { + by_what: Some(ByWhat::Digest(DIRECTORY_A.digest())), + ..Default::default() + })) + .await; + + let mut rx = resp.expect("must succeed").into_inner().into_inner(); + + // The stream should contain one element, an error with Code::NotFound. + let item = rx + .recv() + .await + .expect("must be some") + .expect_err("must be err"); + assert_eq!(item.code(), tonic::Code::NotFound); + + // … and nothing else + assert!(rx.recv().await.is_none()); + + Ok(()) +} + +/// Put a Directory into the store, get it back. +#[tokio::test] +async fn put_get() -> anyhow::Result<()> { + let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?; + + let streaming_request = tonic_mock::streaming_request(vec![DIRECTORY_A.clone()]); + let put_resp = service + .put(streaming_request) + .await + .expect("must succeed") + .into_inner(); + + // the sent root_digest should match the calculated digest + assert_eq!(put_resp.root_digest, DIRECTORY_A.digest()); + + // get it back + let items = get_directories( + &service, + GetDirectoryRequest { + by_what: Some(ByWhat::Digest(DIRECTORY_A.digest().to_vec())), + ..Default::default() + }, + ) + .await + .expect("must not error"); + + assert_eq!(vec![DIRECTORY_A.clone()], items); + + Ok(()) +} + +/// Put multiple Directories into the store, and get them back +#[tokio::test] +async fn put_get_multiple() -> anyhow::Result<()> { + let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?; + + // sending "b" (which refers to "a") without sending "a" first should fail. + let put_resp = service + .put(tonic_mock::streaming_request(vec![DIRECTORY_B.clone()])) + .await + .expect_err("must fail"); + + assert_eq!(tonic::Code::InvalidArgument, put_resp.code()); + + // sending "a", then "b" should succeed, and the response should contain the digest of b. + let put_resp = service + .put(tonic_mock::streaming_request(vec![ + DIRECTORY_A.clone(), + DIRECTORY_B.clone(), + ])) + .await + .expect("must succeed"); + + assert_eq!(DIRECTORY_B.digest(), put_resp.into_inner().root_digest); + + // now, request b, first in non-recursive mode. + let items = get_directories( + &service, + GetDirectoryRequest { + recursive: false, + by_what: Some(ByWhat::Digest(DIRECTORY_B.digest())), + }, + ) + .await + .expect("must not error"); + + // We expect to only get b. + assert_eq!(vec![DIRECTORY_B.clone()], items); + + // now, request b, but in recursive mode. + let items = get_directories( + &service, + GetDirectoryRequest { + recursive: true, + by_what: Some(ByWhat::Digest(DIRECTORY_B.digest())), + }, + ) + .await + .expect("must not error"); + + // We expect to get b, and then a, because that's how we traverse down. + assert_eq!(vec![DIRECTORY_B.clone(), DIRECTORY_A.clone()], items); + + Ok(()) +} |