Skip to content

Commit e931736

Browse files
committed
fixup! Add KvStoreTestSuite.
1 parent 4128e82 commit e931736

File tree

1 file changed

+59
-59
lines changed

1 file changed

+59
-59
lines changed

rust/api/src/kv_store.rs

Lines changed: 59 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,16 @@ use crate::types::{
55
};
66
use async_trait::async_trait;
77

8-
#[cfg(test)]
98
use crate::types::KeyValue;
10-
#[cfg(test)]
119
use bytes::Bytes;
12-
#[cfg(test)]
1310
use rand::distributions::Alphanumeric;
14-
#[cfg(test)]
1511
use rand::{thread_rng, Rng};
1612

17-
pub(crate) const GLOBAL_VERSION_KEY: &str = "global_version";
18-
pub(crate) const INITIAL_RECORD_VERSION: i32 = 1;
13+
/// The key used to store and retrieve the global version of the store.
14+
pub const GLOBAL_VERSION_KEY: &str = "global_version";
15+
16+
/// The initial version number assigned to newly created records.
17+
pub const INITIAL_RECORD_VERSION: i32 = 1;
1918

2019
/// An interface that must be implemented by every backend implementation of VSS.
2120
#[async_trait]
@@ -48,7 +47,6 @@ macro_rules! define_kv_store_tests {
4847
use crate::api::error::VssError;
4948
use crate::api::kv_store::KvStoreTestSuite;
5049
use async_trait::async_trait;
51-
5250
struct $test_suite_name;
5351

5452
#[async_trait]
@@ -89,11 +87,14 @@ macro_rules! define_kv_store_tests {
8987
};
9088
}
9189

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

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

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

409-
let mut previous_page: Option<ListKeyVersionsResponse> = None;
410+
let mut next_page_token: Option<String> = None;
410411
let mut all_key_versions: Vec<KeyValue> = Vec::new();
411412

412413
loop {
413-
let current_page = if previous_page.is_none() {
414-
let page = ctx.list(None, None, None).await?;
415-
assert_eq!(page.global_version, Some(1003));
416-
page
417-
} else {
418-
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
419-
let page = ctx.list(next_page_token, None, None).await?;
420-
assert!(page.global_version.is_none());
421-
page
414+
let current_page = match next_page_token.take() {
415+
None => {
416+
let page = ctx.list(None, None, None).await?;
417+
assert_eq!(page.global_version, Some(1003));
418+
page
419+
},
420+
Some(next_page_token) => {
421+
let page = ctx.list(Some(next_page_token), None, None).await?;
422+
assert!(page.global_version.is_none());
423+
page
424+
},
422425
};
423426

424427
if current_page.key_versions.is_empty() {
425428
break;
426429
}
427430

428-
all_key_versions.extend(current_page.key_versions.clone());
429-
previous_page = Some(current_page);
431+
all_key_versions.extend(current_page.key_versions);
432+
next_page_token = current_page.next_page_token;
430433
}
431434

432-
let unique_keys: std::collections::HashSet<String> =
433-
all_key_versions.iter().map(|kv| kv.key.clone()).collect();
434-
assert_eq!(unique_keys.len(), total_kv_objects as usize);
435-
assert!(!unique_keys.contains(GLOBAL_VERSION_KEY));
436-
437435
if let Some(k1_response) = all_key_versions.iter().find(|kv| kv.key == "k1") {
438436
assert_eq!(k1_response.key, "k1");
439437
assert_eq!(k1_response.version, 2);
@@ -446,6 +444,11 @@ pub(crate) trait KvStoreTestSuite {
446444
assert_eq!(k2_response.value, Bytes::new());
447445
}
448446

447+
let unique_keys: std::collections::HashSet<String> =
448+
all_key_versions.into_iter().map(|kv| kv.key).collect();
449+
assert_eq!(unique_keys.len(), total_kv_objects as usize);
450+
assert!(!unique_keys.contains(GLOBAL_VERSION_KEY));
451+
449452
Ok(())
450453
}
451454

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

462-
let mut previous_page: Option<ListKeyVersionsResponse> = None;
465+
let mut next_page_token: Option<String> = None;
463466
let mut all_key_versions: Vec<KeyValue> = Vec::new();
464467
let key_prefix = "1";
465468

466469
loop {
467-
let current_page = if previous_page.is_none() {
468-
ctx.list(None, Some(page_size), Some(key_prefix.to_string())).await?
469-
} else {
470-
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
471-
ctx.list(next_page_token, Some(page_size), Some(key_prefix.to_string())).await?
470+
let current_page = match next_page_token.take() {
471+
None => ctx.list(None, Some(page_size), Some(key_prefix.to_string())).await?,
472+
Some(next_page_token) => {
473+
ctx.list(Some(next_page_token), Some(page_size), Some(key_prefix.to_string()))
474+
.await?
475+
},
472476
};
473477

474478
if current_page.key_versions.is_empty() {
475479
break;
476480
}
477481

478482
assert!(current_page.key_versions.len() <= page_size as usize);
479-
all_key_versions.extend(current_page.key_versions.clone());
480-
previous_page = Some(current_page);
483+
all_key_versions.extend(current_page.key_versions);
484+
next_page_token = current_page.next_page_token;
481485
}
482486

483487
let unique_keys: std::collections::HashSet<String> =
484-
all_key_versions.iter().map(|kv| kv.key.clone()).collect();
488+
all_key_versions.into_iter().map(|kv| kv.key).collect();
485489

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

507-
let mut previous_page: Option<ListKeyVersionsResponse> = None;
511+
let mut next_page_token: Option<String> = None;
508512
let mut all_key_versions: Vec<KeyValue> = Vec::new();
509513

510514
loop {
511-
let current_page = if previous_page.is_none() {
512-
let page = ctx.list(None, None, None).await?;
513-
assert_eq!(page.global_version.unwrap_or(0), 0);
514-
page
515-
} else {
516-
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
517-
ctx.list(next_page_token, None, None).await?
515+
let current_page = match next_page_token.take() {
516+
None => {
517+
let page = ctx.list(None, None, None).await?;
518+
assert_eq!(page.global_version.unwrap_or(0), 0);
519+
page
520+
},
521+
Some(next_page_token) => ctx.list(Some(next_page_token), None, None).await?,
518522
};
519523

520524
if current_page.key_versions.is_empty() {
521525
break;
522526
}
523527

524-
all_key_versions.extend(current_page.key_versions.clone());
525-
previous_page = Some(current_page);
528+
all_key_versions.extend(current_page.key_versions);
529+
next_page_token = current_page.next_page_token;
526530
}
527531

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

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

546-
let mut previous_page: Option<ListKeyVersionsResponse> = None;
550+
let mut next_page_token: Option<String> = None;
547551
let mut all_key_versions: Vec<KeyValue> = Vec::new();
548552

549553
loop {
550-
let current_page = if previous_page.is_none() {
551-
ctx.list(None, None, None).await?
552-
} else {
553-
let next_page_token = previous_page.as_ref().unwrap().next_page_token.clone();
554-
ctx.list(next_page_token, None, None).await?
554+
let current_page = match next_page_token.take() {
555+
None => ctx.list(None, None, None).await?,
556+
Some(next_page_token) => ctx.list(Some(next_page_token), None, None).await?,
555557
};
556-
557558
if current_page.key_versions.is_empty() {
558559
break;
559560
}
@@ -562,8 +563,8 @@ pub(crate) trait KvStoreTestSuite {
562563
current_page.key_versions.len() < vss_arbitrary_page_size_max as usize,
563564
"Page size exceeds the maximum allowed size"
564565
);
565-
all_key_versions.extend(current_page.key_versions.clone());
566-
previous_page = Some(current_page);
566+
all_key_versions.extend(current_page.key_versions);
567+
next_page_token = current_page.next_page_token;
567568
}
568569

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

575-
#[cfg(test)]
576+
/// Represents the context used for testing [`KvStore`] operations.
576577
pub struct TestContext<'a> {
577578
kv_store: &'a dyn KvStore,
578579
user_token: String,
579580
store_id: String,
580581
}
581582

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

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

0 commit comments

Comments
 (0)