Implement M3.5.7: Rate limiting + idempotency (20 tests)
This commit is contained in:
+251
-95
@@ -1,122 +1,278 @@
|
||||
//! Integration tests for rate limiting and idempotency
|
||||
//!
|
||||
//! Tests rate limiter per apikey, per endpoint, and idempotency caching for ingest.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
#[path = "../crates/mem-cli/src/rate_limiter.rs"]
|
||||
mod rate_limiter;
|
||||
|
||||
#[path = "../crates/mem-cli/src/idempotency.rs"]
|
||||
mod idempotency;
|
||||
|
||||
use rate_limiter::{RateLimiter, LimitConfig};
|
||||
use idempotency::IdempotencyStore;
|
||||
use serde_json::json;
|
||||
|
||||
// ============================================================================
|
||||
// M3.5.7 — Rate limiting per-apikey per-endpoint + idempotency
|
||||
// ============================================================================
|
||||
//
|
||||
// 8 tests covering rate limiting strategy and idempotency
|
||||
//
|
||||
|
||||
/// Test 1: Within limit succeeds
|
||||
#[test]
|
||||
fn r1_rate_limits_per_endpoint() {
|
||||
let limits = json!({
|
||||
"POST /memory/ingest": 100,
|
||||
"GET /memory/query": 1000,
|
||||
"GET /memory/skills": -1,
|
||||
"GET /memory/projects": 100
|
||||
});
|
||||
fn a1_within_limit_succeeds() {
|
||||
let config = LimitConfig {
|
||||
ingest_per_hour: 10.0,
|
||||
query_per_hour: 10.0,
|
||||
projects_per_hour: 10.0,
|
||||
burst_per_second: 10.0,
|
||||
};
|
||||
let limiter = RateLimiter::new(config);
|
||||
|
||||
assert_eq!(limits["POST /memory/ingest"], 100);
|
||||
assert_eq!(limits["GET /memory/query"], 1000);
|
||||
assert_eq!(limits["GET /memory/skills"], -1);
|
||||
// First 10 requests should succeed
|
||||
for i in 0..10 {
|
||||
assert!(
|
||||
limiter.check("user1", "/memory/query").is_ok(),
|
||||
"Request {} should succeed",
|
||||
i + 1
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Test 2: At burst cap, 11th request is rejected
|
||||
#[test]
|
||||
fn r2_rate_limits_per_apikey() {
|
||||
let api_key_1 = "key-abc";
|
||||
let api_key_2 = "key-xyz";
|
||||
|
||||
let mut counters = std::collections::HashMap::new();
|
||||
counters.insert(api_key_1, 5);
|
||||
counters.insert(api_key_2, 2);
|
||||
|
||||
assert_eq!(counters[api_key_1], 5);
|
||||
assert_eq!(counters[api_key_2], 2);
|
||||
fn a2_at_burst_cap_429() {
|
||||
let config = LimitConfig {
|
||||
ingest_per_hour: 100.0,
|
||||
query_per_hour: 10.0,
|
||||
projects_per_hour: 100.0,
|
||||
burst_per_second: 10.0,
|
||||
};
|
||||
let limiter = RateLimiter::new(config);
|
||||
|
||||
// First 10 requests (burst capacity) should succeed
|
||||
for i in 0..10 {
|
||||
let result = limiter.check("user1", "/memory/query");
|
||||
assert!(result.is_ok(), "Request {} within burst should succeed", i + 1);
|
||||
}
|
||||
|
||||
// 11th should fail due to rate limit
|
||||
let result = limiter.check("user1", "/memory/query");
|
||||
assert!(result.is_err(), "11th request should exceed rate limit");
|
||||
|
||||
if let Err(e) = result {
|
||||
assert!(e.retry_after_seconds > 0, "Should have retry_after > 0");
|
||||
}
|
||||
}
|
||||
|
||||
/// Test 3: Limit window (hour) respected
|
||||
#[test]
|
||||
fn r3_burst_allowance() {
|
||||
let burst_capacity: u32 = 10;
|
||||
let requests_in_burst = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
|
||||
|
||||
assert!(requests_in_burst.len() as u32 <= burst_capacity);
|
||||
|
||||
let request_11 = 11;
|
||||
assert!(request_11 as u32 > burst_capacity);
|
||||
fn a3_limit_window_reset() {
|
||||
let config = LimitConfig {
|
||||
ingest_per_hour: 5.0,
|
||||
query_per_hour: 1.0,
|
||||
projects_per_hour: 1.0,
|
||||
burst_per_second: 10.0,
|
||||
};
|
||||
let limiter = RateLimiter::new(config);
|
||||
|
||||
// Use up 5 requests
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check("user1", "/memory/ingest").is_ok());
|
||||
}
|
||||
|
||||
// 6th should fail
|
||||
assert!(limiter.check("user1", "/memory/ingest").is_err());
|
||||
|
||||
// Note: We can't easily test the 1-hour reset without time mocking.
|
||||
// This test documents the behavior but relies on unit tests for time-based refill.
|
||||
}
|
||||
|
||||
/// Test 4: Per-apikey isolation
|
||||
#[test]
|
||||
fn r4_429_response_on_rate_limit() {
|
||||
let response = json!({
|
||||
"status": 429,
|
||||
"error": "rate_limit_exceeded",
|
||||
"reason": "100 requests/hour for POST /memory/ingest",
|
||||
"retry_after_seconds": 47
|
||||
});
|
||||
fn a4_per_apikey_isolation() {
|
||||
let config = LimitConfig {
|
||||
ingest_per_hour: 5.0,
|
||||
query_per_hour: 10.0,
|
||||
projects_per_hour: 10.0,
|
||||
burst_per_second: 10.0,
|
||||
};
|
||||
let limiter = RateLimiter::new(config);
|
||||
|
||||
assert_eq!(response["status"], 429);
|
||||
assert_eq!(response["error"], "rate_limit_exceeded");
|
||||
assert!(response["retry_after_seconds"].is_number());
|
||||
// apikey1 uses up 5 ingest requests
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check("apikey1", "/memory/ingest").is_ok());
|
||||
}
|
||||
assert!(limiter.check("apikey1", "/memory/ingest").is_err());
|
||||
|
||||
// apikey2 should have its own 5 ingest limit
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check("apikey2", "/memory/ingest").is_ok());
|
||||
}
|
||||
assert!(limiter.check("apikey2", "/memory/ingest").is_err());
|
||||
}
|
||||
|
||||
/// Test 5: Per-endpoint isolation
|
||||
#[test]
|
||||
fn r5_retry_after_header() {
|
||||
let retry_after_seconds = 47u32;
|
||||
assert!(retry_after_seconds > 0 && retry_after_seconds <= 3600);
|
||||
fn a5_per_endpoint_isolation() {
|
||||
let config = LimitConfig {
|
||||
ingest_per_hour: 5.0,
|
||||
query_per_hour: 10.0,
|
||||
projects_per_hour: 10.0,
|
||||
burst_per_second: 10.0,
|
||||
};
|
||||
let limiter = RateLimiter::new(config);
|
||||
|
||||
// Use up 5 ingest requests
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check("user1", "/memory/ingest").is_ok());
|
||||
}
|
||||
assert!(limiter.check("user1", "/memory/ingest").is_err());
|
||||
|
||||
// Query endpoint should have separate 10 limit
|
||||
for _ in 0..10 {
|
||||
assert!(limiter.check("user1", "/memory/query").is_ok());
|
||||
}
|
||||
assert!(limiter.check("user1", "/memory/query").is_err());
|
||||
|
||||
// Projects endpoint should have separate 10 limit
|
||||
for _ in 0..10 {
|
||||
assert!(limiter.check("user1", "/memory/projects").is_ok());
|
||||
}
|
||||
assert!(limiter.check("user1", "/memory/projects").is_err());
|
||||
}
|
||||
|
||||
/// Test 6: Retry-After header has correct value
|
||||
#[test]
|
||||
fn r6_idempotency_by_ingest_id() {
|
||||
let ingest_id = "abc123def456abc123def456abc123def456abc123def456abc123def456ab00";
|
||||
|
||||
let job_id_1 = "ingest-uuid-1";
|
||||
let job_id_2 = "ingest-uuid-1";
|
||||
|
||||
assert_eq!(job_id_1, job_id_2);
|
||||
|
||||
let ingest_id_2 = "abc123def456abc123def456abc123def456abc123def456abc123def456ab01";
|
||||
let job_id_3 = "ingest-uuid-2";
|
||||
|
||||
assert_ne!(job_id_1, job_id_3);
|
||||
fn a6_retry_after_header() {
|
||||
let config = LimitConfig {
|
||||
ingest_per_hour: 1.0,
|
||||
query_per_hour: 10.0,
|
||||
projects_per_hour: 10.0,
|
||||
burst_per_second: 10.0,
|
||||
};
|
||||
let limiter = RateLimiter::new(config);
|
||||
|
||||
// Use up limit
|
||||
assert!(limiter.check("user1", "/memory/ingest").is_ok());
|
||||
|
||||
// Next request should fail with retry_after
|
||||
let result = limiter.check("user1", "/memory/ingest");
|
||||
assert!(result.is_err());
|
||||
|
||||
if let Err(e) = result {
|
||||
assert!(e.retry_after_seconds > 0);
|
||||
assert_eq!(e.limit_window_secs, 3600);
|
||||
}
|
||||
}
|
||||
|
||||
/// Test 7: Ingest idempotency — same ingest_id returns cached response
|
||||
#[test]
|
||||
fn r7_idempotency_ttl_24_hours() {
|
||||
let created_at: u64 = 1000000;
|
||||
let queried_at_fresh: u64 = 1000000 + 3600; // 1 hour later (fresh)
|
||||
let queried_at_expired: u64 = 1000000 + (86400 * 2); // 2 days later (expired)
|
||||
|
||||
let ttl_seconds: u64 = 86400; // 24 hours
|
||||
|
||||
let elapsed_fresh = queried_at_fresh - created_at;
|
||||
let elapsed_expired = queried_at_expired - created_at;
|
||||
|
||||
assert!(elapsed_fresh < ttl_seconds, "1 hour should be within TTL");
|
||||
assert!(elapsed_expired > ttl_seconds, "2 days should exceed TTL");
|
||||
fn a7_ingest_id_idempotent() {
|
||||
let store = IdempotencyStore::new(3600);
|
||||
let response1 = json!({"ingest_id": "abc123", "status": "pending", "job_id": "job1"});
|
||||
|
||||
// Store response
|
||||
store.set("abc123".to_string(), response1.clone());
|
||||
|
||||
// Get should return same response
|
||||
let cached = store.get("abc123");
|
||||
assert_eq!(cached, Some(response1));
|
||||
}
|
||||
|
||||
/// Test 8: Different ingest_ids are separate
|
||||
#[test]
|
||||
fn r8_token_bucket_model() {
|
||||
// Token bucket refill model
|
||||
// Capacity: 100 tokens
|
||||
// Refill rate: 100/3600 tokens/sec (100/hour)
|
||||
// Cost per request: 1 token
|
||||
|
||||
let capacity: f32 = 100.0;
|
||||
let refill_rate: f32 = 100.0 / 3600.0; // ~0.0278 tokens/sec
|
||||
let cost_per_request: f32 = 1.0;
|
||||
|
||||
// Simulate tokens over time
|
||||
let mut tokens: f32 = capacity;
|
||||
|
||||
// After 1 hour, bucket refilled
|
||||
let elapsed_1hour: f32 = 3600.0;
|
||||
let refilled_1hour: f32 = (elapsed_1hour * refill_rate).min(capacity);
|
||||
tokens = (tokens + refilled_1hour).min(capacity);
|
||||
|
||||
assert!(tokens >= 50.0 && tokens <= capacity);
|
||||
|
||||
// Make a request (costs 1 token)
|
||||
tokens -= cost_per_request;
|
||||
assert!(tokens < capacity);
|
||||
fn a8_different_ingest_ids_separate() {
|
||||
let store = IdempotencyStore::new(3600);
|
||||
let response_a = json!({"ingest_id": "id_a", "job_id": "job1"});
|
||||
let response_b = json!({"ingest_id": "id_b", "job_id": "job2"});
|
||||
|
||||
store.set("id_a".to_string(), response_a.clone());
|
||||
store.set("id_b".to_string(), response_b.clone());
|
||||
|
||||
assert_eq!(store.get("id_a"), Some(response_a));
|
||||
assert_eq!(store.get("id_b"), Some(response_b));
|
||||
}
|
||||
|
||||
/// Test 9: Idempotency cache expires after TTL
|
||||
#[test]
|
||||
fn a9_idempotency_expires() {
|
||||
let store = IdempotencyStore::new(1); // 1 second TTL
|
||||
let response = json!({"ingest_id": "abc123", "status": "pending"});
|
||||
|
||||
store.set("abc123".to_string(), response);
|
||||
assert!(store.get("abc123").is_some());
|
||||
|
||||
// Wait for expiry
|
||||
std::thread::sleep(Duration::from_millis(1100));
|
||||
|
||||
// Should be expired
|
||||
assert!(store.get("abc123").is_none());
|
||||
}
|
||||
|
||||
/// Test 10: Rate limit configuration documented
|
||||
#[test]
|
||||
fn a10_rate_limit_per_endpoint_documented() {
|
||||
// This test verifies that all endpoints have defined limits
|
||||
let config = LimitConfig::default();
|
||||
|
||||
assert!(config.ingest_per_hour > 0.0, "ingest limit must be > 0");
|
||||
assert!(config.query_per_hour > 0.0, "query limit must be > 0");
|
||||
assert!(config.projects_per_hour > 0.0, "projects limit must be > 0");
|
||||
assert!(config.burst_per_second > 0.0, "burst limit must be > 0");
|
||||
|
||||
// Defaults should be reasonable
|
||||
assert_eq!(config.ingest_per_hour, 100.0); // 100 per hour
|
||||
assert_eq!(config.query_per_hour, 1000.0); // 1000 per hour
|
||||
assert_eq!(config.projects_per_hour, 100.0); // 100 per hour
|
||||
assert_eq!(config.burst_per_second, 10.0); // 10 req/sec burst
|
||||
}
|
||||
|
||||
/// Test 11: Multiple users don't interfere
|
||||
#[test]
|
||||
fn a11_isolation_across_users() {
|
||||
let config = LimitConfig {
|
||||
ingest_per_hour: 3.0,
|
||||
query_per_hour: 3.0,
|
||||
projects_per_hour: 3.0,
|
||||
burst_per_second: 10.0,
|
||||
};
|
||||
let limiter = Arc::new(RateLimiter::new(config));
|
||||
|
||||
// Simulate 3 concurrent users
|
||||
let limiter_a = limiter.clone();
|
||||
let limiter_b = limiter.clone();
|
||||
let limiter_c = limiter.clone();
|
||||
|
||||
// User A: 3 requests
|
||||
for _ in 0..3 {
|
||||
assert!(limiter_a.check("user_a", "/memory/query").is_ok());
|
||||
}
|
||||
|
||||
// User B: 3 requests (independent of user A)
|
||||
for _ in 0..3 {
|
||||
assert!(limiter_b.check("user_b", "/memory/query").is_ok());
|
||||
}
|
||||
|
||||
// User C: 3 requests (independent of A and B)
|
||||
for _ in 0..3 {
|
||||
assert!(limiter_c.check("user_c", "/memory/query").is_ok());
|
||||
}
|
||||
|
||||
// All should be at limit
|
||||
assert!(limiter_a.check("user_a", "/memory/query").is_err());
|
||||
assert!(limiter_b.check("user_b", "/memory/query").is_err());
|
||||
assert!(limiter_c.check("user_c", "/memory/query").is_err());
|
||||
}
|
||||
|
||||
/// Test 12: Idempotency store eviction
|
||||
#[test]
|
||||
fn a12_idempotency_evict_expired() {
|
||||
let store = IdempotencyStore::new(1);
|
||||
store.set("key1".to_string(), json!({"data": "1"}));
|
||||
store.set("key2".to_string(), json!({"data": "2"}));
|
||||
|
||||
std::thread::sleep(Duration::from_millis(1100));
|
||||
store.evict_expired();
|
||||
|
||||
// Both should be gone after eviction
|
||||
assert!(store.get("key1").is_none());
|
||||
assert!(store.get("key2").is_none());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user