use limiteron::postgres_storage::{PostgresStorage, PostgresStorageConfig};
use limiteron::storage::{BanStorage, QuotaStorage};
use std::time::Duration;
const DEFAULT_LIMIT: u64 = 1000;
const DEFAULT_WINDOW: Duration = Duration::from_secs(60);
#[tokio::test]
#[ignore] async fn test_postgres_connection() {
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron");
let storage = PostgresStorage::new(config).await.unwrap();
let result = storage.ping().await;
assert!(result.is_ok());
}
#[tokio::test]
#[ignore]
async fn test_postgres_quota_storage() {
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron");
let storage = PostgresStorage::new(config).await.unwrap();
let user_id = "test_user_pg";
let resource = "test_resource";
let _ = storage
.reset(user_id, resource, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await;
let result = storage
.consume(user_id, resource, 100, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await
.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 900);
let quota = storage.get_quota(user_id, resource).await.unwrap();
assert!(quota.is_some());
assert_eq!(quota.unwrap().consumed, 100);
storage
.reset(user_id, resource, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await
.unwrap();
let quota = storage.get_quota(user_id, resource).await.unwrap();
assert!(quota.is_some());
assert_eq!(quota.unwrap().consumed, 0);
}
#[tokio::test]
#[ignore]
async fn test_postgres_transaction_rollback() {
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron");
let storage = PostgresStorage::new(config).await.unwrap();
let user_id = "transaction_test_user";
let resource = "transaction_test_resource";
let _ = storage
.reset(user_id, resource, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await;
let result1 = storage
.consume(user_id, resource, 100, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await
.unwrap();
assert!(result1.allowed);
storage
.reset(user_id, resource, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await
.unwrap();
let quota = storage.get_quota(user_id, resource).await.unwrap();
assert!(quota.is_some());
assert_eq!(quota.unwrap().consumed, 0);
}
#[tokio::test]
#[ignore]
async fn test_postgres_ban_storage() {
use chrono::Utc;
use limiteron::storage::{BanRecord, BanTarget};
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron");
let storage = PostgresStorage::new(config).await.unwrap();
let target = BanTarget::UserId("test_user_ban".to_string());
let _ = storage.remove_ban(&target).await;
let ban = BanRecord {
target: target.clone(),
ban_times: 1,
duration: Duration::from_secs(60),
banned_at: Utc::now(),
expires_at: Utc::now() + Duration::from_secs(60),
is_manual: false,
reason: "Test ban".to_string(),
};
storage.add_ban(&ban).await.unwrap();
let result = storage.get_ban(&target).await.unwrap();
assert!(result.is_some());
assert_eq!(result.unwrap().ban_times, 1);
storage.remove_ban(&target).await.unwrap();
let result = storage.get_ban(&target).await.unwrap();
assert!(result.is_none());
}
#[tokio::test]
#[ignore]
async fn test_postgres_list_bans() {
use chrono::Utc;
use limiteron::storage::{BanRecord, BanTarget};
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron");
let storage = PostgresStorage::new(config).await.unwrap();
for i in 0..5 {
let target = BanTarget::UserId(format!("list_test_user_{}", i));
let _ = storage.remove_ban(&target).await;
}
for i in 0..5 {
let ban = BanRecord {
target: BanTarget::UserId(format!("list_test_user_{}", i)),
ban_times: 1,
duration: Duration::from_secs(3600),
banned_at: Utc::now(),
expires_at: Utc::now() + Duration::from_secs(3600),
is_manual: false,
reason: "Test ban".to_string(),
};
storage.save(&ban).await.unwrap();
}
for i in 0..5 {
let target = BanTarget::UserId(format!("list_test_user_{}", i));
let history = storage.get_history(&target).await.unwrap();
assert!(history.is_some());
assert_eq!(history.unwrap().ban_times, 1);
}
}
#[tokio::test]
#[ignore]
async fn test_postgres_cleanup_expired_bans() {
use chrono::Utc;
use limiteron::storage::{BanRecord, BanTarget};
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron");
let storage = PostgresStorage::new(config).await.unwrap();
let target = BanTarget::Ip("192.168.1.250".to_string());
let _ = storage.remove_ban(&target).await;
let ban = BanRecord {
target: target.clone(),
ban_times: 1,
duration: Duration::from_secs(1),
banned_at: Utc::now() - Duration::from_secs(10),
expires_at: Utc::now() - Duration::from_secs(5),
is_manual: false,
reason: "Expired ban".to_string(),
};
storage.add_ban(&ban).await.unwrap();
let target2 = BanTarget::Ip("192.168.1.251".to_string());
let _ = storage.remove_ban(&target2).await;
let ban2 = BanRecord {
target: target2.clone(),
ban_times: 1,
duration: Duration::from_secs(3600),
banned_at: Utc::now(),
expires_at: Utc::now() + Duration::from_secs(3600),
is_manual: false,
reason: "Active ban".to_string(),
};
storage.add_ban(&ban2).await.unwrap();
let cleaned = storage.cleanup_expired_bans().await.unwrap();
assert!(cleaned >= 1);
let result = storage.get_ban(&target).await.unwrap();
assert!(result.is_none());
let result = storage.get_ban(&target2).await.unwrap();
assert!(result.is_some());
}
#[tokio::test]
#[ignore]
async fn test_postgres_connection_pool() {
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron")
.with_pool_size(10);
let storage = PostgresStorage::new(config).await.unwrap();
let mut handles = vec![];
for i in 0..20 {
let storage_clone = storage.clone();
let user_id = format!("pool_user_{}", i);
let resource = "pool_resource";
handles.push(tokio::spawn(async move {
for _ in 0..10 {
let _ = storage_clone
.consume(&user_id, resource, 1, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await;
}
}));
}
for handle in handles {
handle.await.unwrap();
}
}
#[tokio::test]
#[ignore]
async fn test_postgres_high_concurrency() {
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron")
.with_pool_size(50);
let storage = PostgresStorage::new(config).await.unwrap();
let user_id = "high_concurrency_pg_user";
let resource = "high_concurrency_pg_resource";
let _ = storage
.reset(user_id, resource, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await;
let mut handles = vec![];
for _ in 0..500 {
let storage_clone = storage.clone();
let user_id = user_id.to_string();
let resource = resource.to_string();
handles.push(tokio::spawn(async move {
storage_clone
.consume(&user_id, &resource, 1, DEFAULT_LIMIT, DEFAULT_WINDOW)
.await
}));
}
let mut success_count = 0;
let mut fail_count = 0;
for handle in handles {
match handle.await.unwrap() {
Ok(result) => {
if result.allowed {
success_count += 1;
} else {
fail_count += 1;
}
}
Err(_) => fail_count += 1,
}
}
println!("Success: {}, Fail: {}", success_count, fail_count);
assert!(success_count + fail_count == 500);
}
#[tokio::test]
#[ignore]
async fn test_postgres_ban_times_tracking() {
use limiteron::storage::BanTarget;
let config =
PostgresStorageConfig::new("postgresql://limiteron:limiteron123@localhost:5432/limiteron");
let storage = PostgresStorage::new(config).await.unwrap();
let target = BanTarget::UserId("ban_times_user".to_string());
let _ = storage.remove_ban(&target).await;
let ban_times = storage.get_ban_times(&target).await.unwrap();
assert_eq!(ban_times, 0);
let new_times = storage.increment_ban_times(&target).await.unwrap();
assert_eq!(new_times, 1);
let new_times = storage.increment_ban_times(&target).await.unwrap();
assert_eq!(new_times, 2);
let ban_times = storage.get_ban_times(&target).await.unwrap();
assert_eq!(ban_times, 2);
}