From 8e727c58e53f820a273cbff31c99b28d63f960bb Mon Sep 17 00:00:00 2001 From: Ion Koutsouris <15728914+ion-elgreco@users.noreply.github.com> Date: Thu, 9 Jan 2025 18:05:24 +0100 Subject: [PATCH] chore: add lakefs client tests --- crates/lakefs/Cargo.toml | 1 + crates/lakefs/src/client.rs | 175 ++++++++++++++++++++++++++++++++++++ 2 files changed, 176 insertions(+) diff --git a/crates/lakefs/Cargo.toml b/crates/lakefs/Cargo.toml index 2058d2775a..c501c0903f 100644 --- a/crates/lakefs/Cargo.toml +++ b/crates/lakefs/Cargo.toml @@ -41,6 +41,7 @@ pretty_env_logger = "0.5.0" rand = "0.8" which = "7" maplit = "1" +mockito = { version = "1.6.1"} [features] integration_test_lakefs = [] \ No newline at end of file diff --git a/crates/lakefs/src/client.rs b/crates/lakefs/src/client.rs index 06cc21f58e..1f229cfd6e 100644 --- a/crates/lakefs/src/client.rs +++ b/crates/lakefs/src/client.rs @@ -276,3 +276,178 @@ impl LakeFSClient { struct LakeFSErrorResponse { message: String, } + +#[cfg(test)] +mod tests { + use std::sync::OnceLock; + + use super::*; + use mockito; + use reqwest::StatusCode; + use tokio::runtime::Runtime; + use uuid::Uuid; + + #[inline] + fn rt() -> &'static Runtime { + static TOKIO_RT: OnceLock = OnceLock::new(); + TOKIO_RT.get_or_init(|| Runtime::new().expect("Failed to create a tokio runtime.")) + } + + #[test] + fn test_create_branch() { + let mut server = mockito::Server::new(); + let mock = server + .mock("POST", "/api/v1/repositories/test_repo/branches") + .with_status(StatusCode::CREATED.as_u16().into()) + .with_body("") + .create(); + + let config = LakeFSConfig::new( + server.url(), + "test_user".to_string(), + "test_pass".to_string(), + ); + let client = LakeFSClient::with_config(config); + let operation_id = Uuid::new_v4(); + let source_url = Url::parse("lakefs://test_repo/main/table").unwrap(); + + let result = rt().block_on(async { client.create_branch(&source_url, operation_id).await }); + assert!(result.is_ok()); + let (new_url, branch_name) = result.unwrap(); + assert_eq!(branch_name, format!("delta-tx-{}", operation_id)); + assert!(new_url.as_str().contains("lakefs://test_repo")); + mock.assert(); + } + + #[test] + fn test_delete_branch() { + let mut server = mockito::Server::new(); + let mock = server + .mock( + "DELETE", + "/api/v1/repositories/test_repo/branches/delta-tx-1234", + ) + .with_status(StatusCode::NO_CONTENT.as_u16().into()) + .create(); + + let config = LakeFSConfig::new( + server.url(), + "test_user".to_string(), + "test_pass".to_string(), + ); + let client = LakeFSClient::with_config(config); + + let result = rt().block_on(async { + client + .delete_branch("test_repo".to_string(), "delta-tx-1234".to_string()) + .await + }); + assert!(result.is_ok()); + mock.assert(); + } + + #[test] + fn test_commit() { + let mut server = mockito::Server::new(); + let mock = server + .mock( + "POST", + "/api/v1/repositories/test_repo/branches/delta-tx-1234/commits", + ) + .with_status(StatusCode::CREATED.as_u16().into()) + .create(); + + let config = LakeFSConfig::new( + server.url(), + "test_user".to_string(), + "test_pass".to_string(), + ); + let client = LakeFSClient::with_config(config); + + let result = rt().block_on(async { + client + .commit( + "test_repo".to_string(), + "delta-tx-1234".to_string(), + "Test commit".to_string(), + false, + ) + .await + }); + assert!(result.is_ok()); + mock.assert(); + } + + #[test] + fn test_merge() { + let mut server = mockito::Server::new(); + let mock = server.mock("POST", "/api/v1/repositories/test_repo/refs/test_transaction_branch/merge/test_target_branch") + .with_status(StatusCode::OK.as_u16().into()) + .create(); + + let config = LakeFSConfig::new( + server.url(), + "test_user".to_string(), + "test_pass".to_string(), + ); + let client = LakeFSClient::with_config(config); + + let result = rt().block_on(async { + client + .merge( + "test_repo".to_string(), + "test_target_branch".to_string(), + "test_transaction_branch".to_string(), + 1, + "Merge commit".to_string(), + false, + ) + .await + }); + assert!(result.is_ok()); + mock.assert(); + } + + #[test] + fn test_decompose_url() { + let config = LakeFSConfig::new( + "http://localhost:8000".to_string(), + "user".to_string(), + "pass".to_string(), + ); + let client = LakeFSClient::with_config(config); + + let (repo, branch, table) = + client.decompose_url("lakefs://test_repo/test_branch/test_table".to_string()); + assert_eq!(repo, "test_repo"); + assert_eq!(branch, "test_branch"); + assert_eq!(table, "test_table"); + + let (repo, branch, table) = + client.decompose_url("lakefs://test_repo/test_branch/data/test_table".to_string()); + assert_eq!(repo, "test_repo"); + assert_eq!(branch, "test_branch"); + assert_eq!(table, "data/test_table"); + } + + #[test] + fn test_transaction_management() { + let config = LakeFSConfig::new( + "http://localhost".to_string(), + "user".to_string(), + "pass".to_string(), + ); + let client = LakeFSClient::with_config(config); + + let transaction_id = Uuid::new_v4(); + let branch_name = "test_branch".to_string(); + + client.set_transaction(transaction_id, branch_name.clone()); + let retrieved_branch = client.get_transaction(transaction_id).unwrap(); + assert_eq!(retrieved_branch, branch_name); + + client.clear_transaction(transaction_id); + let result = client.get_transaction(transaction_id); + assert!(result.is_err()); + } +}