diff --git a/Cargo.toml b/Cargo.toml index 3a09a956..b1ba10fc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/src/rate_limit.rs b/src/rate_limit.rs index 058f9204..ebe7c5bd 100644 --- a/src/rate_limit.rs +++ b/src/rate_limit.rs @@ -160,30 +160,18 @@ impl RateLimiter for TokenLimiter { #[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 = @@ -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 = @@ -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). @@ -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);