Better handling of table setup for async tests

Author: coderdan

src/encrypted_table/mod.rs

@@ -397,17 +397,15 @@ impl EncryptedTable<Dynamo> {
         })
     }
 
+    /// Get a record from the table by primary key from the default dataset.
     pub async fn get<T>(&self, k: impl Into<T::PrimaryKey>) -> Result<Option<T>, GetError>
     where
         T: Decryptable + Identifiable,
     {
-        // TODO: Don't unwrap
-        let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), None)
-            .await
-            .unwrap();
-        self.get_inner(k, scoped_cipher).await
+        self.get_inner(k, None).await
     }
 
+    /// Get a record from the table by primary key from a specific dataset.
     pub async fn get_via<T>(
         &self,
         k: impl Into<T::PrimaryKey>,
@@ -416,21 +414,19 @@ impl EncryptedTable<Dynamo> {
     where
         T: Decryptable + Identifiable,
     {
-        // TODO: Don't unwrap
-        let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), Some(dataset_id))
-            .await
-            .unwrap();
-        self.get_inner(k, scoped_cipher).await
+        self.get_inner(k, Some(dataset_id)).await
     }
 
     async fn get_inner<T>(
         &self,
         k: impl Into<T::PrimaryKey>,
-        cipher: ScopedZeroKmsCipher,
+        dataset_id: Option<Uuid>,
     ) -> Result<Option<T>, GetError>
     where
         T: Decryptable + Identifiable,
     {
+        let cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id).await?;
+
         let PrimaryKeyParts { pk, sk } =
             encrypt_primary_key_parts(&cipher, PreparedPrimaryKey::new::<T>(k))?;
 
@@ -451,13 +447,15 @@ impl EncryptedTable<Dynamo> {
         }
     }
 
+    /// Delete a record from the table by primary key from the default dataset.
     pub async fn delete<E: Searchable + Identifiable>(
         &self,
         k: impl Into<E::PrimaryKey>,
     ) -> Result<(), DeleteError> {
         self.delete_inner::<E>(k.into(), None).await
     }
 
+    /// Delete a record from the table by primary key from a specific dataset.
     pub async fn delete_via<E: Searchable + Identifiable>(
         &self,
         k: impl Into<E::PrimaryKey>,
@@ -489,13 +487,15 @@ impl EncryptedTable<Dynamo> {
         Ok(())
     }
 
+    /// Put a record into the table using the default dataset.
     pub async fn put<T>(&self, record: T) -> Result<(), PutError>
     where
         T: Searchable + Identifiable,
     {
         self.put_inner(record, None).await
     }
 
+    /// Put a record into the table using a specific dataset.
     pub async fn put_via<T>(&self, record: T, dataset_id: Uuid) -> Result<(), PutError>
     where
         T: Searchable + Identifiable,

src/encrypted_table/query.rs

@@ -5,6 +5,7 @@ use cipherstash_client::encryption::{
 };
 use itertools::Itertools;
 use std::{borrow::Cow, collections::HashMap, marker::PhantomData};
+use uuid::Uuid;
 
 use crate::{
     traits::{Decryptable, Searchable},
@@ -132,14 +133,32 @@ impl<S> QueryBuilder<S, &EncryptedTable<Dynamo>>
 where
     S: Searchable + Identifiable,
 {
-    // TODO: Add load_via
+    /// Load all records of type `T` matching the query.
+    /// The default dataset is used.
+    ///
+    /// While a client can decrypt records from any dataset it has access to,
+    /// queries are always scoped to a single dataset.
     pub async fn load<T>(self) -> Result<Vec<T>, QueryError>
     where
         T: Decryptable + Identifiable,
     {
-        let scoped_cipher = ScopedZeroKmsCipher::init(self.storage.cipher.clone(), None)
-            .await
-            .unwrap();
+        self.load_inner(None).await
+    }
+
+    /// Similar to `load`, but the query is scoped to a specific dataset.
+    pub async fn load_via<T>(self, dataset_id: Uuid) -> Result<Vec<T>, QueryError>
+    where
+        T: Decryptable + Identifiable,
+    {
+        self.load_inner(Some(dataset_id)).await
+    }
+
+    async fn load_inner<T>(self, dataset_id: Option<Uuid>) -> Result<Vec<T>, QueryError>
+    where
+        T: Decryptable + Identifiable,
+    {
+        let scoped_cipher =
+            ScopedZeroKmsCipher::init(self.storage.cipher.clone(), dataset_id).await?;
 
         let storage = self.storage;
         let query = self.build()?;

src/errors/mod.rs

@@ -48,6 +48,9 @@ pub enum GetError {
     Encryption(#[from] EncryptionError),
     #[error("AwsError: {0}")]
     Aws(String),
+
+    #[error("ZeroKMS Error: {0}")]
+    ZeroKMS(#[from] zerokms::Error),
 }
 
 /// Error returned by `EncryptedTable::delete` when indexing and deleting records in DynamoDB
@@ -100,6 +103,9 @@ pub enum QueryError {
 
     #[error(transparent)]
     DynamoError(#[from] SdkError<operation::query::QueryError>),
+
+    #[error("ZeroKMS Error: {0}")]
+    ZeroKMS(#[from] zerokms::Error),
 }
 
 pub trait DynamoError: std::error::Error + Sized {}

tests/common.rs

@@ -5,9 +5,91 @@ use aws_sdk_dynamodb::{
     },
     Client,
 };
+use cipherstash_dynamodb::EncryptedTable;
+use miette::Diagnostic;
+use std::{env, future::Future, sync::OnceLock};
+use uuid::Uuid;
 
-pub async fn create_table(client: &Client, table_name: &str) {
+static SECONDARY_DATASET_ID: OnceLock<Uuid> = OnceLock::new();
+
+#[derive(Debug, thiserror::Error, Diagnostic)]
+#[error("Check failed: {0}")]
+pub struct CheckFailed(String);
+
+#[allow(dead_code)]
+pub fn check_eq<A, B>(a: A, b: B) -> miette::Result<()>
+where
+    A: std::fmt::Debug + PartialEq<B>,
+    B: std::fmt::Debug,
+{
+    if a == b {
+        Ok(())
+    } else {
+        Err(CheckFailed(format!("Expected {:?} to equal {:?}", a, b)).into())
+    }
+}
+
+#[allow(dead_code)]
+pub fn check_err<R, E>(result: Result<R, E>) -> miette::Result<()>
+where
+    E: std::fmt::Debug,
+    R: std::fmt::Debug,
+{
+    if result.is_err() {
+        Ok(())
+    } else {
+        Err(CheckFailed(format!("Expected error, got {:?}", result)).into())
+    }
+}
+
+#[allow(dead_code)]
+pub fn check_none<R>(result: Option<R>) -> miette::Result<()>
+where
+    R: std::fmt::Debug,
+{
+    if result.is_none() {
+        Ok(())
+    } else {
+        Err(CheckFailed(format!("Expected None, got {:?}", result)).into())
+    }
+}
+
+#[allow(dead_code)]
+pub fn fail_not_found() -> CheckFailed {
+    CheckFailed("Record not found".into())
+}
+
+/// Run a test with an encrypted table.
+/// The table will be created before the test and deleted after the test.
+/// The name is used as a prefix in case its helpful to distinguish between tests.
+/// A random is appended to the name to ensure uniqueness for async tests.
+#[allow(dead_code)]
+pub async fn with_encrypted_table<F: Future<Output = miette::Result<()>>>(
+    table_name: &str,
+    mut f: impl FnMut(EncryptedTable) -> F,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let config = aws_config::from_env()
+        .endpoint_url("http://localhost:8000")
+        .load()
+        .await;
+
+    let table_name = format!("{}-{}", table_name, Uuid::new_v4());
+    let client = aws_sdk_dynamodb::Client::new(&config);
+
+    create_table(&client, &table_name).await;
+    let table = EncryptedTable::init(client.clone(), &table_name).await?;
+    let result = f(table).await;
+
+    delete_table(&client, &table_name).await;
+    Ok(result?)
+}
+
+pub async fn delete_table(client: &Client, table_name: &str) {
     let _ = client.delete_table().table_name(table_name).send().await;
+}
+
+pub async fn create_table(client: &Client, table_name: &str) {
+    delete_table(client, table_name).await;
 
     client
         .create_table()
@@ -84,6 +166,16 @@ pub async fn create_table(client: &Client, table_name: &str) {
         .expect("Failed to create table");
 }
 
+#[allow(dead_code)]
+pub fn secondary_dataset_id() -> Uuid {
+    *SECONDARY_DATASET_ID.get_or_init(|| {
+        env::var("TEST_SECOND_DATASET_ID")
+            .expect("TEST_SECOND_DATASET_ID must be set")
+            .parse()
+            .expect("TEST_SECOND_DATASET_ID must be a valid UUID")
+    })
+}
+
 #[macro_export]
 macro_rules! assert_err {
     ($cond:expr,) => {

tests/round_trip_tests.rs

@@ -1,8 +1,7 @@
-use cipherstash_client::ZeroKMSConfig;
 use cipherstash_dynamodb::{
-    encrypted_table::Dynamo, Decryptable, Encryptable, EncryptedTable, Identifiable, Searchable,
+    Decryptable, Encryptable, Identifiable, Searchable,
 };
-use miette::IntoDiagnostic;
+use common::{check_eq, check_err, check_none, fail_not_found, secondary_dataset_id, with_encrypted_table};
 use uuid::Uuid;
 mod common;
 
@@ -196,86 +195,50 @@ fn build_test_record(email: &str, name: &str) -> Crazy {
     }
 }
 
-async fn init_table() -> EncryptedTable<Dynamo> {
-    let config = aws_config::from_env()
-        .endpoint_url("http://localhost:8000")
-        .load()
-        .await;
-
-    let client = aws_sdk_dynamodb::Client::new(&config);
-
-    let table_name = "crazy-record";
-
-    common::create_table(&client, table_name).await;
-
-    EncryptedTable::init(client, table_name)
-        .await
-        .expect("Failed to init table")
-}
-
 #[tokio::test]
 async fn test_round_trip() -> Result<(), Box<dyn std::error::Error>> {
-    let table = init_table().await;
-    let record = build_test_record("[email protected]", "Dan");
-    table.put(record.clone()).await.into_diagnostic()?;
-
-    let s: Crazy = table
-        .get(("[email protected]", "Dan"))
-        .await
-        .into_diagnostic()?
-        .unwrap();
-
-    assert_eq!(s, record);
-
-    Ok(())
+    with_encrypted_table("round-trip", |table| async move {
+        let record = build_test_record("[email protected]", "Dan");
+        table.put(record.clone()).await?;
+
+        let s: Crazy = table
+            .get(("[email protected]", "Dan"))
+            .await?
+            .ok_or(fail_not_found())?;
+
+        check_eq(s, record)
+    })
+    .await
 }
 
 #[tokio::test]
 async fn test_invalid_dataset() -> Result<(), Box<dyn std::error::Error>> {
-    let table = init_table().await;
-    let record = build_test_record("[email protected]", "Dan");
+    with_encrypted_table("round-trip", |table| async move {
+        let record = build_test_record("[email protected]", "Dan");
 
-    // A random UUID doesn't exist
-    assert_err!(table.put_via(record.clone(), Uuid::new_v4()).await);
-
-    Ok(())
+        // A random UUID doesn't exist
+        check_err(table.put_via(record.clone(), Uuid::new_v4()).await)
+    })
+    .await
 }
 
 #[tokio::test]
-async fn test_invalid_specific_dataset() -> miette::Result<()> {
-    // TODO: Load client ID from env
-    let client_id = Uuid::parse_str("b91e5b26-f21f-4694-8bce-c61c10e42301").into_diagnostic()?;
-    let client = ZeroKMSConfig::builder()
-        .with_env()
-        .build()
-        .into_diagnostic()?
-        .create_client();
-
-    let dataset = client
-        .create_dataset("test-dataset", "Test dataset")
-        .await
-        .into_diagnostic()?;
-
-    // Grant ourselves access to the dataset
-    client
-        .grant_dataset(client_id, dataset.id)
-        .await
-        .into_diagnostic()?;
-
-    let table = init_table().await;
-    let record = build_test_record("[email protected]", "Person");
-
-    table.put_via(record.clone(), dataset.id).await?;
-
-    let s: Crazy = table
-        .get_via(("[email protected]", "Person"), dataset.id)
-        .await?
-        .unwrap();
-
-    assert_eq!(s, record);
-
-    // Test that we can't get the record via the default dataset
-    assert_none!(table.get::<Crazy>(("[email protected]", "Person")).await?);
-
-    Ok(())
+async fn test_invalid_specific_dataset() -> Result<(), Box<dyn std::error::Error>> {
+    with_encrypted_table("round-trip", |table| async move {
+        let record = build_test_record("[email protected]", "Person");
+        table
+            .put_via(record.clone(), secondary_dataset_id())
+            .await?;
+
+        let s: Crazy = table
+            .get_via(("[email protected]", "Person"), secondary_dataset_id())
+            .await?
+            .ok_or(fail_not_found())?;
+
+        check_eq(s, record)?;
+
+        // Test that we can't get the record via the default dataset
+        check_none(table.get::<Crazy>(("[email protected]", "Person")).await?)
+    })
+    .await
 }