use aingle_cortex::middleware::RateLimiter;
use std::net::{IpAddr, Ipv4Addr};
use std::time::Duration;
use tokio::time::sleep;
#[tokio::test]
async fn test_rate_limiter_basic() {
let limiter = RateLimiter::new(60); let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
let result = limiter.check(ip);
assert!(result.is_ok(), "First request should succeed");
}
#[tokio::test]
async fn test_rate_limiter_exhaustion() {
let limiter = RateLimiter::new(10); let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for i in 0..10 {
let result = limiter.check(ip);
assert!(
result.is_ok(),
"Request {} should succeed (within limit)",
i + 1
);
}
let result = limiter.check(ip);
assert!(result.is_err(), "Request should be rate limited");
}
#[tokio::test]
async fn test_rate_limiter_remaining_count() {
let limiter = RateLimiter::new(100).with_burst_capacity(100);
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
let remaining = limiter.check(ip).expect("Should succeed");
assert_eq!(remaining, 99, "Should have 99 tokens remaining");
for _ in 0..9 {
limiter.check(ip).expect("Should succeed");
}
let remaining = limiter.check(ip).expect("Should succeed");
assert_eq!(remaining, 89, "Should have 89 tokens remaining");
}
#[tokio::test]
async fn test_rate_limiter_multiple_ips() {
let limiter = RateLimiter::new(5);
let ip1 = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
let ip2 = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2));
let ip3 = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
for _ in 0..5 {
limiter.check(ip1).expect("Should succeed");
}
assert!(limiter.check(ip1).is_err(), "IP1 should be rate limited");
assert!(limiter.check(ip2).is_ok(), "IP2 should not be limited");
assert!(limiter.check(ip3).is_ok(), "IP3 should not be limited");
}
#[tokio::test]
async fn test_rate_limiter_token_refill() {
let limiter = RateLimiter::new(60); let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for _ in 0..60 {
limiter.check(ip).expect("Should succeed");
}
assert!(limiter.check(ip).is_err(), "Should be rate limited");
sleep(Duration::from_secs(2)).await;
assert!(
limiter.check(ip).is_ok(),
"Should succeed after token refill"
);
}
#[tokio::test]
async fn test_rate_limiter_burst_capacity() {
let limiter = RateLimiter::new(60).with_burst_capacity(10);
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for i in 0..10 {
assert!(
limiter.check(ip).is_ok(),
"Request {} should succeed",
i + 1
);
}
assert!(limiter.check(ip).is_err(), "Should exceed burst capacity");
}
#[tokio::test]
async fn test_rate_limiter_bucket_info() {
let limiter = RateLimiter::new(100).with_burst_capacity(50);
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
limiter.check(ip).expect("Should succeed");
let (remaining, capacity) = limiter.bucket_info(ip).expect("Bucket should exist");
assert_eq!(remaining, 49, "Should have 49 tokens after first check");
assert_eq!(capacity, 50, "Capacity should be 50");
for _ in 0..10 {
limiter.check(ip).expect("Should succeed");
}
let (remaining, capacity) = limiter.bucket_info(ip).expect("Bucket should exist");
assert_eq!(remaining, 39, "Should have 39 tokens remaining");
assert_eq!(capacity, 50, "Capacity should remain 50");
}
#[tokio::test]
async fn test_rate_limiter_cleanup() {
let limiter = RateLimiter::new(100);
let ip1 = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
let ip2 = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2));
limiter.check(ip1).expect("Should succeed");
limiter.check(ip2).expect("Should succeed");
assert!(limiter.bucket_info(ip1).is_some());
assert!(limiter.bucket_info(ip2).is_some());
limiter.cleanup(Duration::from_secs(0));
}
#[tokio::test]
async fn test_rate_limiter_concurrent_requests() {
use std::sync::Arc;
let limiter = Arc::new(RateLimiter::new(100));
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
let mut handles = vec![];
for _ in 0..50 {
let limiter = limiter.clone();
let handle = tokio::spawn(async move { limiter.check(ip) });
handles.push(handle);
}
let mut successes = 0;
for handle in handles {
if let Ok(Ok(_)) = handle.await {
successes += 1;
}
}
assert_eq!(successes, 50, "All concurrent requests should succeed");
}
#[tokio::test]
async fn test_rate_limiter_ipv6() {
use std::net::Ipv6Addr;
let limiter = RateLimiter::new(10);
let ip = IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1));
for _ in 0..10 {
assert!(limiter.check(ip).is_ok(), "IPv6 should work");
}
assert!(limiter.check(ip).is_err(), "Should be rate limited");
}
#[tokio::test]
async fn test_rate_limiter_different_rates() {
let slow_limiter = RateLimiter::new(10);
let fast_limiter = RateLimiter::new(1000);
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for _ in 0..10 {
slow_limiter.check(ip).expect("Should succeed");
}
assert!(slow_limiter.check(ip).is_err());
for _ in 0..100 {
fast_limiter.check(ip).expect("Should succeed");
}
}
#[tokio::test]
async fn test_rate_limiter_retry_after() {
let limiter = RateLimiter::new(60); let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for _ in 0..60 {
limiter.check(ip).expect("Should succeed");
}
let err = limiter.check(ip).expect_err("Should be rate limited");
let err_string = err.to_string();
assert!(
err_string.contains("Retry after"),
"Error should contain retry-after info"
);
}
#[tokio::test]
async fn test_rate_limiter_gradual_refill() {
let limiter = RateLimiter::new(60); let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for _ in 0..60 {
limiter.check(ip).expect("Should succeed");
}
assert!(limiter.check(ip).is_err(), "Should be rate limited");
sleep(Duration::from_secs(3)).await;
let mut successes = 0;
for _ in 0..5 {
if limiter.check(ip).is_ok() {
successes += 1;
} else {
break;
}
}
assert!(
successes >= 2 && successes <= 4,
"Should refill approximately 3 tokens, got {}",
successes
);
}
#[tokio::test]
async fn test_rate_limiter_max_capacity() {
let limiter = RateLimiter::new(100).with_burst_capacity(20);
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for _ in 0..20 {
limiter.check(ip).expect("Should succeed");
}
assert!(limiter.check(ip).is_err());
sleep(Duration::from_secs(30)).await;
let mut successes = 0;
for _ in 0..25 {
if limiter.check(ip).is_ok() {
successes += 1;
} else {
break;
}
}
assert!(
successes <= 21, "Should not exceed capacity, got {}",
successes
);
}
#[test]
fn test_rate_limiter_default() {
let limiter = RateLimiter::default();
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for _ in 0..100 {
assert!(limiter.check(ip).is_ok());
}
assert!(limiter.check(ip).is_err());
}
#[test]
fn test_rate_limiter_builder() {
let limiter = RateLimiter::new(60)
.with_burst_capacity(30)
.with_secure_ip(true);
let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
for _ in 0..30 {
assert!(limiter.check(ip).is_ok());
}
assert!(limiter.check(ip).is_err());
}