Skip to content

Commit

Permalink
fixup! Add KvStoreTestSuite.
Browse files Browse the repository at this point in the history
  • Loading branch information
@G8XSU
G8XSU committed Nov 20, 2024
1 parent 4128e82 commit e931736
Showing 1 changed file with 59 additions and 59 deletions.
118 changes: 59 additions & 59 deletions rust/api/src/kv_store.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,16 @@ use crate::types::{
};
use async_trait::async_trait;

#[cfg(test)]
use crate::types::KeyValue;
#[cfg(test)]
use bytes::Bytes;
#[cfg(test)]
use rand::distributions::Alphanumeric;
#[cfg(test)]
use rand::{thread_rng, Rng};

pub(crate) const GLOBAL_VERSION_KEY: &str = "global_version";
pub(crate) const INITIAL_RECORD_VERSION: i32 = 1;
/// The key used to store and retrieve the global version of the store.
pub const GLOBAL_VERSION_KEY: &str = "global_version";

/// The initial version number assigned to newly created records.
pub const INITIAL_RECORD_VERSION: i32 = 1;

/// An interface that must be implemented by every backend implementation of VSS.
#[async_trait]
Expand Down Expand Up @@ -48,7 +47,6 @@ macro_rules! define_kv_store_tests {
use crate::api::error::VssError;
use crate::api::kv_store::KvStoreTestSuite;
use async_trait::async_trait;

struct $test_suite_name;

#[async_trait]
Expand Down Expand Up @@ -89,11 +87,14 @@ macro_rules! define_kv_store_tests {
};
}

/// Contains tests for a [`KvStore`] implementation to ensure it complies with the VSS protocol.
#[allow(missing_docs)]
#[async_trait]
#[cfg(test)]
pub(crate) trait KvStoreTestSuite {
pub trait KvStoreTestSuite {
/// The type of store being tested. This must implement the [`KvStore`] trait.
type Store: KvStore + 'static;

/// Creates and returns a new instance of the store to be tested.
async fn create_store() -> Self::Store;

async fn put_should_succeed_when_single_object_put_operation() -> Result<(), VssError> {
Expand Down Expand Up @@ -406,34 +407,31 @@ pub(crate) trait KvStoreTestSuite {
ctx.put_objects(Some(1001), vec![kv("k2", "k2v2", 1)]).await?;
ctx.put_objects(Some(1002), vec![kv("k2", "k2v3", 2)]).await?;

let mut previous_page: Option<ListKeyVersionsResponse> = None;
let mut next_page_token: Option<String> = None;
let mut all_key_versions: Vec<KeyValue> = Vec::new();

loop {
let current_page = if previous_page.is_none() {
let page = ctx.list(None, None, None).await?;
assert_eq!(page.global_version, Some(1003));
page
} else {
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
let page = ctx.list(next_page_token, None, None).await?;
assert!(page.global_version.is_none());
page
let current_page = match next_page_token.take() {
None => {
let page = ctx.list(None, None, None).await?;
assert_eq!(page.global_version, Some(1003));
page
},
Some(next_page_token) => {
let page = ctx.list(Some(next_page_token), None, None).await?;
assert!(page.global_version.is_none());
page
},
};

if current_page.key_versions.is_empty() {
break;
}

all_key_versions.extend(current_page.key_versions.clone());
previous_page = Some(current_page);
all_key_versions.extend(current_page.key_versions);
next_page_token = current_page.next_page_token;
}

let unique_keys: std::collections::HashSet<String> =
all_key_versions.iter().map(|kv| kv.key.clone()).collect();
assert_eq!(unique_keys.len(), total_kv_objects as usize);
assert!(!unique_keys.contains(GLOBAL_VERSION_KEY));

if let Some(k1_response) = all_key_versions.iter().find(|kv| kv.key == "k1") {
assert_eq!(k1_response.key, "k1");
assert_eq!(k1_response.version, 2);
Expand All @@ -446,6 +444,11 @@ pub(crate) trait KvStoreTestSuite {
assert_eq!(k2_response.value, Bytes::new());
}

let unique_keys: std::collections::HashSet<String> =
all_key_versions.into_iter().map(|kv| kv.key).collect();
assert_eq!(unique_keys.len(), total_kv_objects as usize);
assert!(!unique_keys.contains(GLOBAL_VERSION_KEY));

Ok(())
}

Expand All @@ -459,34 +462,35 @@ pub(crate) trait KvStoreTestSuite {
ctx.put_objects(Some(i as i64), vec![kv(&format!("{}k", i), "k1v1", 0)]).await?;
}

let mut previous_page: Option<ListKeyVersionsResponse> = None;
let mut next_page_token: Option<String> = None;
let mut all_key_versions: Vec<KeyValue> = Vec::new();
let key_prefix = "1";

loop {
let current_page = if previous_page.is_none() {
ctx.list(None, Some(page_size), Some(key_prefix.to_string())).await?
} else {
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
ctx.list(next_page_token, Some(page_size), Some(key_prefix.to_string())).await?
let current_page = match next_page_token.take() {
None => ctx.list(None, Some(page_size), Some(key_prefix.to_string())).await?,
Some(next_page_token) => {
ctx.list(Some(next_page_token), Some(page_size), Some(key_prefix.to_string()))
.await?
},
};

if current_page.key_versions.is_empty() {
break;
}

assert!(current_page.key_versions.len() <= page_size as usize);
all_key_versions.extend(current_page.key_versions.clone());
previous_page = Some(current_page);
all_key_versions.extend(current_page.key_versions);
next_page_token = current_page.next_page_token;
}

let unique_keys: std::collections::HashSet<String> =
all_key_versions.iter().map(|kv| kv.key.clone()).collect();
all_key_versions.into_iter().map(|kv| kv.key).collect();

assert_eq!(unique_keys.len(), 11);
let expected_keys: std::collections::HashSet<String> =
["1k", "10k", "11k", "12k", "13k", "14k", "15k", "16k", "17k", "18k", "19k"]
.iter()
.into_iter()
.map(|s| s.to_string())
.collect();
assert_eq!(unique_keys, expected_keys);
Expand All @@ -504,29 +508,29 @@ pub(crate) trait KvStoreTestSuite {
ctx.put_objects(None, vec![kv(&format!("k{}", i), "k1v1", 0)]).await?;
}

let mut previous_page: Option<ListKeyVersionsResponse> = None;
let mut next_page_token: Option<String> = None;
let mut all_key_versions: Vec<KeyValue> = Vec::new();

loop {
let current_page = if previous_page.is_none() {
let page = ctx.list(None, None, None).await?;
assert_eq!(page.global_version.unwrap_or(0), 0);
page
} else {
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
ctx.list(next_page_token, None, None).await?
let current_page = match next_page_token.take() {
None => {
let page = ctx.list(None, None, None).await?;
assert_eq!(page.global_version.unwrap_or(0), 0);
page
},
Some(next_page_token) => ctx.list(Some(next_page_token), None, None).await?,
};

if current_page.key_versions.is_empty() {
break;
}

all_key_versions.extend(current_page.key_versions.clone());
previous_page = Some(current_page);
all_key_versions.extend(current_page.key_versions);
next_page_token = current_page.next_page_token;
}

let unique_keys: std::collections::HashSet<String> =
all_key_versions.iter().map(|kv| kv.key.clone()).collect();
all_key_versions.into_iter().map(|kv| kv.key).collect();
assert_eq!(unique_keys.len(), total_kv_objects as usize);
assert!(!unique_keys.contains(GLOBAL_VERSION_KEY));

Expand All @@ -543,17 +547,14 @@ pub(crate) trait KvStoreTestSuite {
ctx.put_objects(Some(i as i64), vec![kv(&format!("k{}", i), "k1v1", 0)]).await?;
}

let mut previous_page: Option<ListKeyVersionsResponse> = None;
let mut next_page_token: Option<String> = None;
let mut all_key_versions: Vec<KeyValue> = Vec::new();

loop {
let current_page = if previous_page.is_none() {
ctx.list(None, None, None).await?
} else {
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
ctx.list(next_page_token, None, None).await?
let current_page = match next_page_token.take() {
None => ctx.list(None, None, None).await?,
Some(next_page_token) => ctx.list(Some(next_page_token), None, None).await?,
};

if current_page.key_versions.is_empty() {
break;
}
Expand All @@ -562,8 +563,8 @@ pub(crate) trait KvStoreTestSuite {
current_page.key_versions.len() < vss_arbitrary_page_size_max as usize,
"Page size exceeds the maximum allowed size"
);
all_key_versions.extend(current_page.key_versions.clone());
previous_page = Some(current_page);
all_key_versions.extend(current_page.key_versions);
next_page_token = current_page.next_page_token;
}

assert_eq!(all_key_versions.len(), total_kv_objects as usize);
Expand All @@ -572,15 +573,15 @@ pub(crate) trait KvStoreTestSuite {
}
}

#[cfg(test)]
/// Represents the context used for testing [`KvStore`] operations.
pub struct TestContext<'a> {
kv_store: &'a dyn KvStore,
user_token: String,
store_id: String,
}

#[cfg(test)]
impl<'a> TestContext<'a> {
/// Creates a new [`TestContext`] with the given [`KvStore`] implementation.
pub fn new(kv_store: &'a dyn KvStore) -> Self {
let store_id: String = (0..7).map(|_| thread_rng().sample(Alphanumeric) as char).collect();
TestContext { kv_store, user_token: "userToken".to_string(), store_id }
Expand Down Expand Up @@ -640,7 +641,6 @@ impl<'a> TestContext<'a> {
}
}

#[cfg(test)]
fn kv(key: &str, value: &str, version: i64) -> KeyValue {
KeyValue { key: key.to_string(), version, value: Bytes::from(value.to_string()) }
}

0 comments on commit e931736

Please sign in to comment.