Skip to content

Commit

Permalink
Use tokio's testing feature in tests for rate limitter instead of moc…
Browse files Browse the repository at this point in the history
…king of TokioClock.
  • Loading branch information
kannapoix committed Sep 3, 2024
1 parent 1cf2906 commit adac81a
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 31 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ rand_core = "0.6.4"
log = "0.4.17"
log4rs = { version = "1.2.0", features = ["file_appender"] }
rcgen = { version = "0.13.1", features = ["pem", "x509-parser"] }
tokio = { version = "1.25.0", features = ["rt", "rt-multi-thread", "signal"] }
tokio = { version = "1.25.0", features = ["rt", "rt-multi-thread", "signal", "test-util"] }
tonic = { version = "0.11", features = [ "tls", "transport" ] }
tonic_lnd = { git = "https://github.com/orbitalturtle/tonic_lnd", rev="18c5a71084886024a6b90307bfb8822288c5daea", package="fedimint-tonic-lnd", features = ["lightningrpc", "routerrpc", "versionrpc"] }
hex = "0.4.3"
Expand Down
42 changes: 12 additions & 30 deletions src/rate_limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,30 +160,18 @@ impl<C: Clock> RateLimiter for TokenLimiter<C> {
#[cfg(test)]
mod tests {
use super::*;
use crate::tests::test_utils::pubkey;
use core::ops::SubAssign;
use mockall::mock;
use tokio::time::{Duration, Instant};
use crate::{clock::TokioClock, tests::test_utils::pubkey};
use tokio::time::Duration;

const TEST_COUNT: u8 = 2;
const TEST_FREQUENCY: Duration = Duration::from_secs(1);

mock! {
FixedClock{}

impl Clock for FixedClock{
fn now(&self) -> Instant;
}
}

#[test]
fn test_peer_connection() {
#[tokio::test(start_paused = true)]
async fn test_peer_connection() {
let pk_0 = pubkey(0);
let pk_1 = pubkey(1);

let mut clock = MockFixedClock::new();
// TODO: use constant value for mocking.
clock.expect_now().returning(|| Instant::now());
let clock = TokioClock::new();

// Assert that we're set up with our original peer.
let mut rate_limiter =
Expand All @@ -199,16 +187,12 @@ mod tests {
assert_eq!(rate_limiter.peers(), vec![pk_1]);
}

#[test]
fn test_rate_limiting() {
#[tokio::test(start_paused = true)]
async fn test_rate_limiting() {
let pk_0 = pubkey(0);
let pk_1 = pubkey(1);

let mut clock = MockFixedClock::new();

// TODO: use constant value for mocking.
let start_time = Instant::now();
clock.expect_now().returning(move || start_time);
let clock = TokioClock::new();

// Assert that we're set up with our original peer.
let mut rate_limiter =
Expand Down Expand Up @@ -238,7 +222,7 @@ mod tests {
rate_limiter.peer_disconnected(pk_1);

// Update our clock to a time which reflects that we need an update.
rate_limiter.last_update.sub_assign(TEST_FREQUENCY);
tokio::time::advance(TEST_FREQUENCY).await;

// The disconnected peer should not be allowed any queries (they're currently unknown to
// us).
Expand All @@ -260,13 +244,11 @@ mod tests {
assert!(!rate_limiter.query_peer(pk_1));
}

#[test]
fn test_query_peer_unknown_peer() {
#[tokio::test(start_paused = true)]
async fn test_query_peer_unknown_peer() {
let pk_0 = pubkey(0);

let mut clock = MockFixedClock::new();
let start_time = Instant::now();
clock.expect_now().returning(move || start_time);
let clock = TokioClock::new();

let mut rate_limiter =
TokenLimiter::new(vec![].into_iter(), TEST_COUNT, TEST_FREQUENCY, clock);
Expand Down

0 comments on commit adac81a

Please sign in to comment.