Skip to content

Commit

Permalink
throttle: add basic test for Redis throttling
Browse files Browse the repository at this point in the history
  • Loading branch information
djc committed Sep 19, 2024
1 parent b55564e commit 1cce908
Showing 1 changed file with 48 additions and 0 deletions.
48 changes: 48 additions & 0 deletions crates/throttle/src/throttle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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_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> {
gcra_throttle(&self.0, key, limit, period, max_burst, quantity).await
}
}

async fn test_big_limits(
limit: u64,
max_burst: Option<u64>,
Expand Down Expand Up @@ -315,4 +347,20 @@ 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_throttle_1_000() {
if !RedisServer::is_available() {
return;
}

let redis = RedisServer::spawn("").await.unwrap();
let conn = redis.connection().await.unwrap();
let cx = dbg!(RedisContext::try_from(conn).await.unwrap());

match cx.cl_throttle_command {
true => test_big_limits(1_000, None, 0.02, &RedisWithCell(cx.connection)).await,
false => test_big_limits(1_000, None, 0.02, &VanillaRedis(cx.connection)).await,
}
}
}

0 comments on commit 1cce908

Please sign in to comment.