use super::{QuotaPolicy, QuotaPool};
use crate::refill_rate::RefillRate;
use std::sync::{Arc, Barrier};
#[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_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);
}