From 31e4e15ebbdf4accb2effda33ce4b1bba8595d1b Mon Sep 17 00:00:00 2001 From: Joe Neeman Date: Fri, 23 Dec 2022 13:52:34 -0600 Subject: [PATCH] Add bigtable emulator and first tests. --- src/bigtable/emulator.rs | 110 ++++++++++++++++++++++++++++ src/bigtable/mod.rs | 4 + tests/bigtable_client.rs | 153 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 267 insertions(+) create mode 100644 src/bigtable/emulator.rs create mode 100644 tests/bigtable_client.rs diff --git a/src/bigtable/emulator.rs b/src/bigtable/emulator.rs new file mode 100644 index 0000000..3f3f9cc --- /dev/null +++ b/src/bigtable/emulator.rs @@ -0,0 +1,110 @@ +//! Testing infra to make use of the bigtable emulator. +//! +//! +//! Follow installation directions from link above to set up your local development. Once setup, +//! you should be able to run the pubsub emulator driven tests. + +use futures::{future::BoxFuture, FutureExt}; + +use crate::{ + bigtable, + builder::ClientBuilder, + emulator::{self, EmulatorData}, +}; + +type BoxError = Box; + +/// Struct to hold a started PubSub emulator process. Process is closed when struct is dropped. +pub struct EmulatorClient { + inner: crate::emulator::EmulatorClient, + instance: String, +} + +const DATA: EmulatorData = EmulatorData { + gcloud_param: "bigtable", + kill_pattern: "bigtable", + availability_check: create_bigtable_client, + extra_args: Vec::new(), +}; + +const INSTANCE_ID: &str = "test-instance"; + +impl EmulatorClient { + /// Create a new emulator instance with a default project name and instance name + pub async fn new() -> Result { + Ok(EmulatorClient { + inner: emulator::EmulatorClient::new(DATA).await?, + instance: INSTANCE_ID.into(), + }) + } + + /// Create a new emulator instance with the given project name and instance name + pub async fn with_project_and_instance( + project_name: impl Into, + instance_name: impl Into, + ) -> Result { + Ok(EmulatorClient { + inner: emulator::EmulatorClient::with_project(DATA, project_name).await?, + instance: instance_name.into(), + }) + } + + /// Get the endpoint at which the emulator is listening for requests + pub fn endpoint(&self) -> String { + self.inner.endpoint() + } + + /// Get the project name with which the emulator was initialized + pub fn project(&self) -> &str { + self.inner.project() + } + + /// Get the instance name with which the emulator was initialized + pub fn instance(&self) -> &str { + &self.instance + } + + /// Get a client builder which is pre-configured to work with this emulator instance + pub fn builder(&self) -> &ClientBuilder { + self.inner.builder() + } + + /// Create a new table under this emulator's given project name and instance name. + /// + /// The column families will be created wth effectively no garbage collection. + pub async fn create_table( + &self, + table_name: &str, + column_families: impl IntoIterator>, + ) -> Result<(), BoxError> { + let config = bigtable::admin::BigtableTableAdminConfig { + endpoint: self.endpoint(), + ..bigtable::admin::BigtableTableAdminConfig::default() + }; + + let mut admin = self + .builder() + .build_bigtable_admin_client(config, &self.project(), &self.instance) + .await?; + + let column_families = column_families + .into_iter() + .map(|name| (name.into(), bigtable::admin::Rule::MaxNumVersions(i32::MAX))); + admin.create_table(table_name, column_families).await?; + + Ok(()) + } +} + +fn create_bigtable_client(port: &str) -> BoxFuture> { + async move { + bigtable::api::bigtable::v2::bigtable_client::BigtableClient::connect(format!( + "http://{}:{}", + crate::emulator::HOST, + port + )) + .await?; + Ok(()) + } + .boxed() +} diff --git a/src/bigtable/mod.rs b/src/bigtable/mod.rs index a48be14..b2eb8ad 100644 --- a/src/bigtable/mod.rs +++ b/src/bigtable/mod.rs @@ -23,6 +23,10 @@ pub mod mutation; pub use client_builder::BigtableConfig; pub use mutation::{MutateRowRequest, MutateRowsError, MutateRowsRequest}; +#[cfg(feature = "emulators")] +#[cfg_attr(docsrs, doc(cfg(feature = "emulators")))] +pub mod emulator; + #[allow(rustdoc::broken_intra_doc_links, rustdoc::bare_urls, missing_docs)] pub mod api { pub mod rpc { diff --git a/tests/bigtable_client.rs b/tests/bigtable_client.rs new file mode 100644 index 0000000..cbbfa20 --- /dev/null +++ b/tests/bigtable_client.rs @@ -0,0 +1,153 @@ +#[cfg(all(feature = "bigtable", feature = "emulators"))] +mod bigtable_client_tests { + use futures::TryStreamExt; + use hyper::body::Bytes; + use ya_gcp::bigtable::{self, admin::Rule, api, emulator::EmulatorClient, ReadRowsRequest}; + + #[tokio::test] + async fn create_table() { + let table_name = "test-table"; + let emulator = EmulatorClient::new().await.unwrap(); + let config = bigtable::admin::BigtableTableAdminConfig::new().endpoint(emulator.endpoint()); + let mut admin = emulator + .builder() + .build_bigtable_admin_client(config, emulator.project(), emulator.instance()) + .await + .unwrap(); + + admin + .create_table(table_name, [("column".into(), Rule::MaxNumVersions(1))]) + .await + .unwrap(); + + let tables: Vec<_> = admin + .list_tables() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + assert_eq!(tables.len(), 1); + assert_eq!( + tables[0].name, + "projects/test-project/instances/test-instance/tables/test-table" + ); + } + + async fn default_client(table_name: &str) -> (EmulatorClient, bigtable::BigtableClient) { + let emulator = EmulatorClient::new().await.unwrap(); + emulator + .create_table(table_name, ["fam1", "fam2"]) + .await + .unwrap(); + + let config = bigtable::BigtableConfig::new().endpoint(emulator.endpoint()); + let client = emulator + .builder() + .build_bigtable_client(config, emulator.project(), emulator.instance()) + .await + .unwrap(); + (emulator, client) + } + + #[tokio::test] + async fn set_and_read_row() { + let table_name = "test-table"; + let (_emulator, mut client) = default_client(table_name).await; + + client + .set_row_data( + table_name, + "fam1".into(), + "row1-key", + [("col1", "data1"), ("col2", "data2")], + ) + .await + .unwrap(); + + let row = client + .read_one_row(table_name, "row1-key") + .await + .unwrap() + .unwrap(); + assert_eq!("row1-key", row.key); + assert_eq!(1, row.families.len()); + assert_eq!("fam1", row.families[0].name); + dbg!(&row.families); + assert_eq!(2, row.families[0].columns.len()); + assert_eq!("col1", row.families[0].columns[0].qualifier); + assert_eq!("data1", row.families[0].columns[0].cells[0].value); + assert_eq!("col2", row.families[0].columns[1].qualifier); + assert_eq!("data2", row.families[0].columns[1].cells[0].value); + + let row_range: Vec<_> = client + .read_row_range( + table_name, + Bytes::from("row1-key")..=Bytes::from("row1-key"), + None, + ) + .try_collect() + .await + .unwrap(); + assert_eq!(row_range, vec![row]); + } + + #[tokio::test] + async fn cell_versions() { + let table_name = "test-table"; + let (emulator, mut client) = default_client(table_name).await; + client + .set_row_data_with_timestamp( + table_name, + "fam1".into(), + // The bigtable emulator enforces millisecond granularity + 6000, + "row1-key", + [("col1", "data1")], + ) + .await + .unwrap(); + client + .set_row_data_with_timestamp( + table_name, + "fam1".into(), + 7000, + "row1-key", + [("col1", "data2")], + ) + .await + .unwrap(); + let row = client + .read_one_row(table_name, "row1-key") + .await + .unwrap() + .unwrap(); + // read_one_row only returns the latest version + assert_eq!(1, row.families[0].columns[0].cells.len()); + assert_eq!("data2", row.families[0].columns[0].cells[0].value); + assert_eq!( + vec!["data2"], + row.most_recent_cells().map(|c| c.value).collect::>() + ); + + let req = ReadRowsRequest { + table_name: format!( + "projects/{}/instances/{}/tables/{table_name}", + emulator.project(), + emulator.instance() + ), + rows: Some(api::bigtable::v2::RowSet::default().with_key("row1-key")), + ..Default::default() + }; + let rows: Vec<_> = client.read_rows(req).try_collect().await.unwrap(); + assert_eq!(1, rows.len()); + let row = &rows[0]; + dbg!(row); + assert_eq!(2, row.families[0].columns[0].cells.len()); + assert_eq!( + vec!["data2"], + row.most_recent_cells().map(|c| c.value).collect::>() + ); + } +}