throttle: add basic tests for Redis throttling

KumoCorp · Sep 19, 2024 · b5e5b76 · b5e5b76
1 parent fa7a3e2
commit b5e5b76
Showing 1 changed file with 60 additions and 0 deletions.
diff --git a/crates/throttle/src/throttle.rs b/crates/throttle/src/throttle.rs
@@ -222,6 +222,8 @@ pub async fn throttle(
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate::redis::RedisContext;
+    use mod_redis::test::RedisServer;
 
     trait Throttler {
         async fn throttle(
@@ -247,6 +249,36 @@ mod test {
         }
     }
 
+    struct RedisWithCell(RedisConnection);
+
+    impl Throttler for RedisWithCell {
+        async fn throttle(
+            &self,
+            key: &str,
+            limit: u64,
+            period: Duration,
+            max_burst: u64,
+            quantity: Option<u64>,
+        ) -> Result<ThrottleResult, Error> {
+            redis_cell_throttle(&self.0, key, limit, period, max_burst, quantity).await
+        }
+    }
+
+    struct VanillaRedis(RedisConnection);
+
+    impl Throttler for VanillaRedis {
+        async fn throttle(
+            &self,
+            key: &str,
+            limit: u64,
+            period: Duration,
+            max_burst: u64,
+            quantity: Option<u64>,
+        ) -> Result<ThrottleResult, Error> {
+            redis_script_throttle(&self.0, key, limit, period, max_burst, quantity).await
+        }
+    }
+
     async fn test_big_limits(
         limit: u64,
         max_burst: Option<u64>,
@@ -315,4 +347,32 @@ mod test {
         // time periods produced by the overally limit, rather than the burst.
         test_big_limits(60_000, Some(30_000), 0.05, &*MEMORY).await;
     }
+
+    #[tokio::test]
+    async fn redis_cell_throttle_1_000() {
+        if !RedisServer::is_available() {
+            return;
+        }
+
+        let redis = RedisServer::spawn("").await.unwrap();
+        let conn = redis.connection().await.unwrap();
+        let cx = RedisContext::try_from(conn).await.unwrap();
+        if !cx.has_redis_cell {
+            return;
+        }
+
+        test_big_limits(1_000, None, 0.02, &RedisWithCell(cx.connection)).await;
+    }
+
+    #[tokio::test]
+    async fn redis_script_throttle_1_000() {
+        if !RedisServer::is_available() {
+            return;
+        }
+
+        let redis = RedisServer::spawn("").await.unwrap();
+        let conn = redis.connection().await.unwrap();
+        let cx = RedisContext::try_from(conn).await.unwrap();
+        test_big_limits(1_000, None, 0.02, &VanillaRedis(cx.connection)).await;
+    }
 }