Merge pull request #168 from kannapoix/improve-rate-limit-testing

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"

src/rate_limit.rs

@@ -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 =
@@ -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);