Skip to content

Commit

Permalink
add test for rate limiter with multiple actions
Browse files Browse the repository at this point in the history
  • Loading branch information
pietroalbini committed Jul 26, 2023
1 parent aeb1356 commit 6635d95
Showing 1 changed file with 83 additions and 0 deletions.
83 changes: 83 additions & 0 deletions src/rate_limiter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -403,6 +403,49 @@ mod tests {
Ok(())
}

#[test]
fn two_actions_dont_interfere_with_each_other() -> QueryResult<()> {
let conn = &mut pg_connection();
let now = now();

let mut config = HashMap::new();
config.insert(
LimitedAction::PublishNew,
RateLimiterConfig {
rate: Duration::from_secs(1),
burst: 10,
},
);
config.insert(
LimitedAction::YankUnyank,
RateLimiterConfig {
rate: Duration::from_secs(1),
burst: 20,
},
);
let rate = RateLimiter::new(config);

let user_id = new_user(conn, "user")?;

assert_eq!(
10,
rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?
.tokens
);
assert_eq!(
9,
rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?
.tokens
);
assert_eq!(
20,
rate.take_token(user_id, LimitedAction::YankUnyank, now, conn)?
.tokens
);

Ok(())
}

#[test]
fn override_is_used_instead_of_global_burst_if_present() -> QueryResult<()> {
let conn = &mut pg_connection();
Expand Down Expand Up @@ -480,6 +523,46 @@ mod tests {
Ok(())
}

#[test]
fn override_is_different_for_each_action() -> QueryResult<()> {
let conn = &mut pg_connection();
let now = now();
let user_id = new_user(conn, "user")?;

let mut config = HashMap::new();
for action in [LimitedAction::PublishNew, LimitedAction::YankUnyank] {
config.insert(
action,
RateLimiterConfig {
rate: Duration::from_secs(1),
burst: 10,
},
);
}
let rate = RateLimiter::new(config);

diesel::insert_into(publish_rate_overrides::table)
.values((
publish_rate_overrides::user_id.eq(user_id),
publish_rate_overrides::action.eq(LimitedAction::PublishNew),
publish_rate_overrides::burst.eq(20),
))
.execute(conn)?;

assert_eq!(
20,
rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?
.tokens,
);
assert_eq!(
10,
rate.take_token(user_id, LimitedAction::YankUnyank, now, conn)?
.tokens,
);

Ok(())
}

fn new_user(conn: &mut PgConnection, gh_login: &str) -> QueryResult<i32> {
use crate::models::NewUser;

Expand Down

0 comments on commit 6635d95

Please sign in to comment.