use crate::prelude::*;
use std::time::{Duration, SystemTime};
const LIMIT_COUNT: usize = 5;
const LIMIT_DURATION: Duration = Duration::from_secs(10);
const LIMIT_DURATION_SHORT: Duration = Duration::from_millis(250);
#[tokio::test]
async fn test_get_wait_duration_empty() {
let limiter = RateLimiter::new(LIMIT_COUNT, LIMIT_DURATION);
let wait = limiter.get_wait_duration().await;
print_duration("Wait", wait);
assert!(wait.is_none());
}
#[tokio::test]
async fn test_get_wait_duration_available() {
let limiter = RateLimiter::new(LIMIT_COUNT, LIMIT_DURATION);
let now = SystemTime::now();
{
let mut requests = limiter.requests.lock().await;
for _ in 0..(LIMIT_COUNT - 1) {
requests.push_back(now);
}
}
let wait = limiter.get_wait_duration().await;
print_duration("Wait", wait);
assert!(wait.is_none());
}
#[tokio::test]
async fn test_get_wait_duration_full() {
let limiter = RateLimiter::new(LIMIT_COUNT, LIMIT_DURATION);
let now = SystemTime::now();
{
let mut requests = limiter.requests.lock().await;
for _ in 0..LIMIT_COUNT {
requests.push_back(now);
}
}
let wait = limiter.get_wait_duration().await;
print_duration("Wait", wait);
assert!(wait.is_some());
assert!(approximately_equals(
wait.expect("full queue should require wait"),
LIMIT_DURATION,
Duration::from_millis(100)
));
}
#[tokio::test]
async fn test_execute_available() {
let limiter = RateLimiter::new(LIMIT_COUNT, LIMIT_DURATION);
let now = SystemTime::now();
{
let mut requests = limiter.requests.lock().await;
for _ in 0..(LIMIT_COUNT - 1) {
requests.push_back(now);
}
}
let now = SystemTime::now();
let wait = limiter.execute().await;
let elapsed = now.elapsed().expect("elapsed should not fail");
print_duration("Wait", wait);
print_duration("Elapsed", Some(elapsed));
assert!(wait.is_none());
assert!(approximately_equals(
elapsed,
Duration::from_secs(0),
Duration::from_millis(50)
));
}
#[tokio::test]
async fn test_execute_full() {
let limiter = RateLimiter::new(LIMIT_COUNT, LIMIT_DURATION_SHORT);
let now = SystemTime::now();
{
let mut requests = limiter.requests.lock().await;
for _ in 0..LIMIT_COUNT {
requests.push_back(now);
}
}
let now = SystemTime::now();
let wait = limiter.execute().await;
let elapsed = now.elapsed().expect("elapsed should not fail");
print_duration("Wait", wait);
print_duration("Elapsed", Some(elapsed));
assert!(wait.is_some());
assert!(approximately_equals(
wait.expect("full queue should require wait"),
LIMIT_DURATION_SHORT,
Duration::from_millis(50)
));
assert!(approximately_equals(
elapsed,
LIMIT_DURATION_SHORT,
Duration::from_millis(50)
));
}
fn approximately_equals(d1: Duration, d2: Duration, tolerance: Duration) -> bool {
if d1 > d2 {
d1.checked_sub(d2)
.expect("subtraction should not underflow")
<= tolerance
} else {
d2.checked_sub(d1)
.expect("subtraction should not underflow")
<= tolerance
}
}
fn print_duration(name: &str, duration: Option<Duration>) {
if let Some(duration) = duration {
println!("{name} duration: {:.3} seconds", duration.as_secs_f64());
} else {
println!("{name} duration: None");
}
}
#[tokio::test]
async fn test_constructor_initializes_empty_queue() {
let limiter = RateLimiter::new(10, Duration::from_secs(60));
assert!(limiter.requests.lock().await.is_empty());
assert_eq!(limiter.rate.num, 10);
assert_eq!(limiter.rate.per, Duration::from_secs(60));
}
#[tokio::test]
async fn test_remove_stale_clears_old_requests() {
let limiter = RateLimiter::new(LIMIT_COUNT, Duration::from_millis(50));
let now = SystemTime::now();
let old_time = now - Duration::from_millis(100);
{
let mut requests = limiter.requests.lock().await;
for _ in 0..LIMIT_COUNT {
requests.push_back(old_time);
}
assert_eq!(requests.len(), LIMIT_COUNT);
}
let wait = limiter.get_wait_duration().await;
assert!(wait.is_none());
assert!(limiter.requests.lock().await.is_empty());
}
#[tokio::test]
async fn test_partial_stale_removal() {
let limiter = RateLimiter::new(LIMIT_COUNT, Duration::from_millis(100));
let now = SystemTime::now();
let stale_time = now - Duration::from_millis(200);
{
let mut requests = limiter.requests.lock().await;
for _ in 0..3 {
requests.push_back(stale_time);
}
for _ in 0..2 {
requests.push_back(now);
}
assert_eq!(requests.len(), 5);
}
let wait = limiter.get_wait_duration().await;
assert!(wait.is_none());
assert_eq!(limiter.requests.lock().await.len(), 2);
}
#[tokio::test]
async fn test_execute_adds_request_to_queue() {
let limiter = RateLimiter::new(LIMIT_COUNT, LIMIT_DURATION);
assert!(limiter.requests.lock().await.is_empty());
limiter.execute().await;
assert_eq!(limiter.requests.lock().await.len(), 1);
}
#[tokio::test]
async fn test_multiple_executes_fill_queue() {
let limiter = RateLimiter::new(LIMIT_COUNT, LIMIT_DURATION);
for _ in 0..(LIMIT_COUNT - 1) {
limiter.execute().await;
}
assert_eq!(limiter.requests.lock().await.len(), LIMIT_COUNT - 1);
let wait = limiter.execute().await;
assert!(wait.is_none());
assert_eq!(limiter.requests.lock().await.len(), LIMIT_COUNT);
}