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, SymlinkNode}; 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() }; static ref DIRECTORY_C: Directory = Directory { directories: vec![ DirectoryNode { name: "a".to_string(), digest: DIRECTORY_A.digest(), size: DIRECTORY_A.size(), }, 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(()) } /// Put multiple Directories into the store, and omit duplicates. #[tokio::test] async fn put_get_dedup() -> anyhow::Result<()> { let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?; // Send "A", then "C", which refers to "A" two times // Pretend we're a dumb client sending A twice. let put_resp = service .put(tonic_mock::streaming_request(vec![ DIRECTORY_A.clone(), DIRECTORY_A.clone(), DIRECTORY_C.clone(), ])) .await .expect("must succeed"); assert_eq!(DIRECTORY_C.digest(), put_resp.into_inner().root_digest); // Ask for "C" recursively. We expect to only get "A" once, as there's no point sending it twice. let items = get_directories( &service, GetDirectoryRequest { recursive: true, by_what: Some(ByWhat::Digest(DIRECTORY_C.digest())), }, ) .await .expect("must not error"); // We expect to get C, and then A (once, as the second A has been deduplicated). assert_eq!(vec![DIRECTORY_C.clone(), DIRECTORY_A.clone()], items); Ok(()) } /// Trying to upload a Directory failing validation should fail. #[tokio::test] async fn put_reject_failed_validation() -> anyhow::Result<()> { let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?; // construct a broken Directory message that fails validation let broken_directory = Directory { symlinks: vec![SymlinkNode { name: "".to_string(), target: "doesntmatter".to_string(), }], ..Default::default() }; assert!(broken_directory.validate().is_err()); // send it over, it must fail let put_resp = service .put(tonic_mock::streaming_request(vec![broken_directory])) .await .expect_err("must fail"); assert_eq!(put_resp.code(), tonic::Code::InvalidArgument); Ok(()) } /// Trying to upload a Directory with wrong size should fail. #[tokio::test] async fn put_reject_wrong_size() -> anyhow::Result<()> { let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?; // Construct a directory referring to DIRECTORY_A, but with wrong size. let broken_parent_directory = Directory { directories: vec![DirectoryNode { name: "foo".to_string(), digest: DIRECTORY_A.digest(), size: 42, }], ..Default::default() }; // Make sure we got the size wrong. assert_ne!( broken_parent_directory.directories[0].size, DIRECTORY_A.size() ); // now upload both (first A, then the broken parent). This must fail. let put_resp = service .put(tonic_mock::streaming_request(vec![ DIRECTORY_A.clone(), broken_parent_directory, ])) .await .expect_err("must fail"); assert_eq!(put_resp.code(), tonic::Code::InvalidArgument); Ok(()) }