use super::{QuotaPolicy, QuotaPool, QuotaSlot};
use crate::refill_rate::RefillRate;
use std::sync::{Arc, Barrier};
use std::time::Duration;
#[test]
fn pool_consumes_initial_tokens_for_one_key() {
let policy = QuotaPolicy::new()
.set_capacity(10.0)
.set_refill_rate(RefillRate::per_sec(3));
let pool = QuotaPool::new(policy, 10);
let mut results = vec![];
for _ in 0..100 {
results.push(pool.consume("testing", 1));
}
assert_eq!(results.iter().filter(|r| r.is_ok()).count(), 10);
assert_eq!(results.iter().filter(|r| r.is_err()).count(), 90);
}
#[test]
fn pool_limits_one_key_across_threads() {
let policy = QuotaPolicy::new().set_capacity(10.0);
let pool = Arc::new(QuotaPool::new(policy, 10));
let barrier = Arc::new(Barrier::new(100));
let mut handles = Vec::new();
for _ in 0..100 {
let pool = Arc::clone(&pool);
let barrier = Arc::clone(&barrier);
handles.push(std::thread::spawn(move || {
barrier.wait();
pool.consume("testing", 1).is_ok()
}));
}
let accepted = handles
.into_iter()
.map(|handle| handle.join().expect("worker thread panicked"))
.filter(|accepted| *accepted)
.count();
assert_eq!(accepted, 10);
assert_eq!(pool.len(), 1);
}
#[test]
fn pool_consumes_owned_key_and_keeps_entry() {
let policy = QuotaPolicy::new().set_capacity(1.0);
let pool = QuotaPool::new(policy, 1);
assert!(pool.consume_owned(String::from("testing"), 1).is_ok());
assert!(pool.consume("testing", 1).is_err());
assert!(pool.contains_key("testing"));
}
#[test]
fn pool_removes_key_and_recreates_fresh_quota() {
let policy = QuotaPolicy::new().set_capacity(1.0);
let pool = QuotaPool::new(policy, 1);
assert!(pool.consume("testing", 1).is_ok());
assert!(pool.consume("testing", 1).is_err());
assert!(pool.remove("testing"));
assert!(!pool.contains_key("testing"));
assert_eq!(pool.len(), 0);
assert!(pool.consume("testing", 1).is_ok());
}
#[test]
fn pool_cache_does_not_use_removed_slot_after_reuse() {
let policy = QuotaPolicy::new().set_capacity(1.0);
let pool = QuotaPool::with_capacity(policy, 1, 2);
assert!(pool.consume("a", 1).is_ok());
assert!(pool.remove("a"));
pool.insert_key("b");
assert!(pool.consume("a", 1).is_ok());
assert!(pool.consume("b", 1).is_ok());
}
#[test]
fn pool_tracks_spread_keys_independently() {
let policy = QuotaPolicy::new().set_capacity(1.0);
let keys: Vec<_> = (0..64).map(|id| format!("key-{id}")).collect();
let pool = Arc::new(QuotaPool::with_capacity(policy, 1, keys.len()));
pool.insert_keys(&keys);
assert_eq!(pool.len(), keys.len());
let mut handles = Vec::new();
for key in keys {
let pool = Arc::clone(&pool);
handles.push(std::thread::spawn(move || {
pool.consume(key.as_str(), 1).is_ok()
}));
}
let accepted = handles
.into_iter()
.map(|handle| handle.join().expect("worker thread panicked"))
.filter(|accepted| *accepted)
.count();
assert_eq!(accepted, 64);
assert_eq!(pool.len(), 64);
}
#[test]
fn slot_skips_refill_before_interval() {
let policy = QuotaPolicy::new()
.set_capacity(10.0)
.set_refill_rate(RefillRate::per_sec(10))
.set_refill_interval(Duration::from_secs(1));
let slot = QuotaSlot::new(&policy, 0);
let result = slot.consume(&policy, 500_000_000, 1);
assert!(result.is_err());
assert_eq!(result.unwrap_err().current().available(), 0.0);
}
#[test]
fn slot_refills_after_interval() {
let policy = QuotaPolicy::new()
.set_capacity(72.0)
.set_refill_rate(RefillRate::per_sec(10))
.set_refill_interval(Duration::from_secs(1));
let slot = QuotaSlot::new(&policy, 0);
assert!(slot.consume(&policy, 500_000_000, 1).is_err());
let result = slot
.consume(&policy, 1_000_000_000, 1)
.expect("interval refill should add tokens");
assert_eq!(result.available(), 9.0);
}
#[test]
fn slot_keeps_initial_tokens_when_interval_skips_refill() {
let policy = QuotaPolicy::new()
.set_capacity(1.0)
.set_refill_rate(RefillRate::per_sec(10))
.set_refill_interval(Duration::from_secs(1));
let slot = QuotaSlot::new(&policy, 10);
let result = slot
.consume(&policy, 500_000_000, 5)
.expect("skipped refill should not clamp existing quota");
assert_eq!(result.available(), 5.0);
}