Skip to content

GitLab

Unverified Commit 1f10ba55 authored by Yan Ru Pei's avatar Yan Ru Pei Committed by GitHub
Browse files

refactor(kv-router): tidy indexer test query helpers (#7801) 


Signed-off-by: default avatarPeaBrane <yanrpei@gmail.com>
parent 3d8e85e2
Show whitespace changes
Inline Side-by-side
Showing with 73 additions and 76 deletions
lib/kv-router/src/indexer/tests.rs
View file @ 1f10ba55
......@@ -252,6 +252,23 @@ async fn flush_and_settle(index: &dyn KvIndexerInterface) {
tokio::time::sleep(Duration::from_millis(100)).await;
}
async fn query_scores(index: &dyn KvIndexerInterface, query: &[u64]) -> OverlapScores {
index
.find_matches(query.iter().copied().map(LocalBlockHash).collect())
.await
.unwrap()
}
async fn assert_score(
index: &dyn KvIndexerInterface,
query: &[u64],
worker: WorkerWithDpRank,
expected_score: u32,
) {
let scores = query_scores(index, query).await;
assert_eq!(scores.scores.get(&worker), Some(&expected_score));
}
async fn assert_query_score_and_tree_size(
index: &dyn KvIndexerInterface,
query: &[u64],
......@@ -259,15 +276,28 @@ async fn assert_query_score_and_tree_size(
expected_score: u32,
expected_tree_size: usize,
) {
let scores = index
.find_matches(query.iter().copied().map(LocalBlockHash).collect())
.await
.unwrap();
let scores = query_scores(index, query).await;
assert_eq!(scores.scores.get(&worker), Some(&expected_score));
assert_eq!(scores.tree_sizes.get(&worker), Some(&expected_tree_size));
}
async fn assert_no_scores(index: &dyn KvIndexerInterface, query: &[u64]) {
let scores = query_scores(index, query).await;
assert!(scores.scores.is_empty());
}
async fn assert_exact_scores(
index: &dyn KvIndexerInterface,
query: &[u64],
expected_scores: &[(WorkerWithDpRank, u32)],
) {
let scores = query_scores(index, query).await;
assert_eq!(scores.scores.len(), expected_scores.len());
for (worker, expected_score) in expected_scores {
assert_eq!(scores.scores.get(worker), Some(expected_score));
}
}
mod interface_tests {
use super::*;
use rstest_reuse::apply;
......@@ -282,17 +312,7 @@ mod interface_tests {
flush_and_settle(index.as_ref()).await;
// Find matches using local hashes
let scores = index
.find_matches(vec![
LocalBlockHash(1),
LocalBlockHash(2),
LocalBlockHash(3),
])
.await
.unwrap();
assert_eq!(scores.scores.len(), 1);
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(0, 0)).unwrap(), 3);
assert_score(index.as_ref(), &[1, 2, 3], WorkerWithDpRank::new(0, 0), 3).await;
}
#[tokio::test]
......@@ -404,16 +424,7 @@ mod interface_tests {
flush_and_settle(index.as_ref()).await;
// Find matches for [1, 2, 999] - should match first 2 then stop
let scores = index
.find_matches(vec![
LocalBlockHash(1),
LocalBlockHash(2),
LocalBlockHash(999),
])
.await
.unwrap();
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(0, 0)).unwrap(), 2);
assert_score(index.as_ref(), &[1, 2, 999], WorkerWithDpRank::new(0, 0), 2).await;
}
#[tokio::test]
......@@ -429,16 +440,7 @@ mod interface_tests {
flush_and_settle(index.as_ref()).await;
// Find should return nothing
let scores = index
.find_matches(vec![
LocalBlockHash(1),
LocalBlockHash(2),
LocalBlockHash(3),
])
.await
.unwrap();
assert!(scores.scores.is_empty());
assert_no_scores(index.as_ref(), &[1, 2, 3]).await;
}
#[tokio::test]
......@@ -454,20 +456,25 @@ mod interface_tests {
flush_and_settle(index.as_ref()).await;
// Query [1] - both workers should match
let scores = index.find_matches(vec![LocalBlockHash(1)]).await.unwrap();
assert_eq!(scores.scores.len(), 2);
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(0, 0)).unwrap(), 1);
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(1, 0)).unwrap(), 1);
assert_exact_scores(
index.as_ref(),
&[1],
&[
(WorkerWithDpRank::new(0, 0), 1),
(WorkerWithDpRank::new(1, 0), 1),
],
)
.await;
// Query [1, 2] - worker 0 matches both, worker 1 matches only first block
let scores = index
.find_matches(vec![LocalBlockHash(1), LocalBlockHash(2)])
.await
.unwrap();
assert_eq!(scores.scores.len(), 2);
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(0, 0)).unwrap(), 2);
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(1, 0)).unwrap(), 1);
assert_exact_scores(
index.as_ref(),
&[1, 2],
&[
(WorkerWithDpRank::new(0, 0), 2),
(WorkerWithDpRank::new(1, 0), 1),
],
)
.await;
}
#[tokio::test]
......@@ -486,16 +493,12 @@ mod interface_tests {
// Allow time for async remove_worker processing
flush_and_settle(index.as_ref()).await;
let scores = index
.find_matches(vec![
LocalBlockHash(1),
LocalBlockHash(2),
LocalBlockHash(3),
])
.await
.unwrap();
assert_eq!(scores.scores.len(), 1);
assert!(scores.scores.contains_key(&WorkerWithDpRank::new(1, 0)));
assert_exact_scores(
index.as_ref(),
&[1, 2, 3],
&[(WorkerWithDpRank::new(1, 0), 3)],
)
.await;
}
#[tokio::test]
......@@ -588,9 +591,7 @@ mod interface_tests {
flush_and_settle(index.as_ref()).await;
// Empty query should return empty scores
let scores = index.find_matches(vec![]).await.unwrap();
assert!(scores.scores.is_empty());
assert_no_scores(index.as_ref(), &[]).await;
}
#[tokio::test]
......@@ -602,12 +603,7 @@ mod interface_tests {
flush_and_settle(index.as_ref()).await;
// Query for non-existent blocks
let scores = index
.find_matches(vec![LocalBlockHash(999), LocalBlockHash(998)])
.await
.unwrap();
assert!(scores.scores.is_empty());
assert_no_scores(index.as_ref(), &[999, 998]).await;
}
#[tokio::test]
......@@ -692,15 +688,16 @@ mod interface_tests {
flush_and_settle(index.as_ref()).await;
// Query for full sequence [1, 2, 3, 4, 5] should match all 5 blocks
let full_seq: Vec<LocalBlockHash> = (1..=5).map(LocalBlockHash).collect();
let scores = index.find_matches(full_seq).await.unwrap();
assert_eq!(scores.scores.len(), 1);
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(0, 0)).unwrap(), 5);
assert_score(
index.as_ref(),
&[1, 2, 3, 4, 5],
WorkerWithDpRank::new(0, 0),
5,
)
.await;
// Query for just [1, 2, 3] should match 3 blocks
let prefix_seq: Vec<LocalBlockHash> = (1..=3).map(LocalBlockHash).collect();
let scores = index.find_matches(prefix_seq).await.unwrap();
assert_eq!(*scores.scores.get(&WorkerWithDpRank::new(0, 0)).unwrap(), 3);
assert_score(index.as_ref(), &[1, 2, 3], WorkerWithDpRank::new(0, 0), 3).await;
}
#[tokio::test]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment