2026-08-26 13:35:50 -07:00
|
|
|
//! 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;
|
2026-08-23 17:19:42 -07:00
|
|
|
use serde_json::json;
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 1: Within limit succeeds
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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);
|
2026-08-23 17:19:42 -07:00
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
// First 10 requests should succeed
|
|
|
|
|
for i in 0..10 {
|
|
|
|
|
assert!(
|
|
|
|
|
limiter.check("user1", "/memory/query").is_ok(),
|
|
|
|
|
"Request {} should succeed",
|
|
|
|
|
i + 1
|
|
|
|
|
);
|
|
|
|
|
}
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 2: At burst cap, 11th request is rejected
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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");
|
|
|
|
|
}
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 3: Limit window (hour) respected
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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.
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 4: Per-apikey isolation
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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);
|
2026-08-23 17:19:42 -07:00
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
// 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());
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 5: Per-endpoint isolation
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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());
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 6: Retry-After header has correct value
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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);
|
|
|
|
|
}
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 7: Ingest idempotency — same ingest_id returns cached response
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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));
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|
|
|
|
|
|
2026-08-26 13:35:50 -07:00
|
|
|
/// Test 8: Different ingest_ids are separate
|
2026-08-23 17:19:42 -07:00
|
|
|
#[test]
|
2026-08-26 13:35:50 -07:00
|
|
|
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());
|
2026-08-23 17:19:42 -07:00
|
|
|
}
|