Skip to content

Commit c138d0c

Browse files
committed
add test for rate limiter with multiple actions
1 parent 6315414 commit c138d0c

File tree

1 file changed

+83
-0
lines changed

1 file changed

+83
-0
lines changed

src/rate_limiter.rs

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,49 @@ mod tests {
403403
Ok(())
404404
}
405405

406+
#[test]
407+
fn two_actions_dont_interfere_with_each_other() -> QueryResult<()> {
408+
let conn = &mut pg_connection();
409+
let now = now();
410+
411+
let mut config = HashMap::new();
412+
config.insert(
413+
LimitedAction::PublishNew,
414+
RateLimiterConfig {
415+
rate: Duration::from_secs(1),
416+
burst: 10,
417+
},
418+
);
419+
config.insert(
420+
LimitedAction::YankUnyank,
421+
RateLimiterConfig {
422+
rate: Duration::from_secs(1),
423+
burst: 20,
424+
},
425+
);
426+
let rate = RateLimiter::new(config);
427+
428+
let user_id = new_user(conn, "user")?;
429+
430+
assert_eq!(
431+
10,
432+
rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?
433+
.tokens
434+
);
435+
assert_eq!(
436+
9,
437+
rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?
438+
.tokens
439+
);
440+
assert_eq!(
441+
20,
442+
rate.take_token(user_id, LimitedAction::YankUnyank, now, conn)?
443+
.tokens
444+
);
445+
446+
Ok(())
447+
}
448+
406449
#[test]
407450
fn override_is_used_instead_of_global_burst_if_present() -> QueryResult<()> {
408451
let conn = &mut pg_connection();
@@ -480,6 +523,46 @@ mod tests {
480523
Ok(())
481524
}
482525

526+
#[test]
527+
fn override_is_different_for_each_action() -> QueryResult<()> {
528+
let conn = &mut pg_connection();
529+
let now = now();
530+
let user_id = new_user(conn, "user")?;
531+
532+
let mut config = HashMap::new();
533+
for action in [LimitedAction::PublishNew, LimitedAction::YankUnyank] {
534+
config.insert(
535+
action,
536+
RateLimiterConfig {
537+
rate: Duration::from_secs(1),
538+
burst: 10,
539+
},
540+
);
541+
}
542+
let rate = RateLimiter::new(config);
543+
544+
diesel::insert_into(publish_rate_overrides::table)
545+
.values((
546+
publish_rate_overrides::user_id.eq(user_id),
547+
publish_rate_overrides::action.eq(LimitedAction::PublishNew),
548+
publish_rate_overrides::burst.eq(20),
549+
))
550+
.execute(conn)?;
551+
552+
assert_eq!(
553+
20,
554+
rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?
555+
.tokens,
556+
);
557+
assert_eq!(
558+
10,
559+
rate.take_token(user_id, LimitedAction::YankUnyank, now, conn)?
560+
.tokens,
561+
);
562+
563+
Ok(())
564+
}
565+
483566
fn new_user(conn: &mut PgConnection, gh_login: &str) -> QueryResult<i32> {
484567
use crate::models::NewUser;
485568

0 commit comments

Comments
 (0)