Skip to content

Commit 3e86592

Browse files
committed
default_policy: LA tests: allow LA customisation
It is needed for at least one test to differ from the default test configuration in some option. To make this possible, the function returning the policy used for tests now accepts a mutator that configures LatencyAwarenessBuilder before DefaultPolicy is built.
1 parent 1a913a1 commit 3e86592

File tree

1 file changed

+17
-5
lines changed

1 file changed

+17
-5
lines changed

scylla/src/transport/load_balancing/default.rs

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1962,9 +1962,9 @@ mod latency_awareness {
19621962
}
19631963
}
19641964

1965-
pub fn latency_aware_default_policy() -> DefaultPolicy {
1966-
let latency_awareness = LatencyAwareness::builder().build();
1967-
1965+
fn default_policy_with_given_latency_awareness(
1966+
latency_awareness: LatencyAwareness,
1967+
) -> DefaultPolicy {
19681968
let pick_predicate = {
19691969
let latency_predicate = latency_awareness.generate_predicate();
19701970
Box::new(move |node: &NodeRef| {
@@ -1983,6 +1983,17 @@ mod latency_awareness {
19831983
}
19841984
}
19851985

1986+
fn latency_aware_default_policy_customised(
1987+
configurer: impl FnOnce(LatencyAwarenessBuilder) -> LatencyAwarenessBuilder,
1988+
) -> DefaultPolicy {
1989+
let latency_awareness = configurer(LatencyAwareness::builder()).build();
1990+
default_policy_with_given_latency_awareness(latency_awareness)
1991+
}
1992+
1993+
fn latency_aware_default_policy() -> DefaultPolicy {
1994+
latency_aware_default_policy_customised(|b| b)
1995+
}
1996+
19861997
#[tokio::test]
19871998
async fn latency_aware_default_policy_does_not_penalise_if_no_latency_info_available_yet() {
19881999
let policy = latency_aware_default_policy();
@@ -2124,8 +2135,9 @@ mod latency_awareness {
21242135

21252136
#[tokio::test]
21262137
async fn latency_aware_default_policy_does_not_penalise_if_retry_period_expired() {
2127-
let mut policy = latency_aware_default_policy();
2128-
policy.latency_awareness.as_mut().unwrap().retry_period = Duration::from_millis(10);
2138+
let policy = latency_aware_default_policy_customised(|b| {
2139+
b.retry_period(Duration::from_millis(10))
2140+
});
21292141

21302142
let cluster = tests::mock_cluster_data_for_token_unaware_tests().await;
21312143

0 commit comments

Comments
 (0)