telemt/src/ip_tracker/tests.rs

use super::*;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::atomic::Ordering;

fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr {
    IpAddr::V4(Ipv4Addr::new(oct1, oct2, oct3, oct4))
}

fn test_ipv6() -> IpAddr {
    IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1))
}

#[tokio::test]
async fn test_basic_ip_limit() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 2).await;

    let ip1 = test_ipv4(192, 168, 1, 1);
    let ip2 = test_ipv4(192, 168, 1, 2);
    let ip3 = test_ipv4(192, 168, 1, 3);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip3).await.is_err());

    assert_eq!(tracker.get_active_ip_count("test_user").await, 2);
}

#[tokio::test]
async fn test_active_window_rejects_new_ip_and_keeps_existing_session() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 1).await;
    tracker
        .set_limit_policy(UserMaxUniqueIpsMode::ActiveWindow, 30)
        .await;

    let ip1 = test_ipv4(10, 10, 10, 1);
    let ip2 = test_ipv4(10, 10, 10, 2);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.is_ip_active("test_user", ip1).await);
    assert!(tracker.check_and_add("test_user", ip2).await.is_err());

    // Existing session remains active; only new unique IP is denied.
    assert!(tracker.is_ip_active("test_user", ip1).await);
    assert_eq!(tracker.get_active_ip_count("test_user").await, 1);
}

#[tokio::test]
async fn test_reconnection_from_same_ip() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 2).await;

    let ip1 = test_ipv4(192, 168, 1, 1);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert_eq!(tracker.get_active_ip_count("test_user").await, 1);
}

#[tokio::test]
async fn test_same_ip_disconnect_keeps_active_while_other_session_alive() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 2).await;

    let ip1 = test_ipv4(192, 168, 1, 1);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert_eq!(tracker.get_active_ip_count("test_user").await, 1);

    tracker.remove_ip("test_user", ip1).await;
    assert_eq!(tracker.get_active_ip_count("test_user").await, 1);

    tracker.remove_ip("test_user", ip1).await;
    assert_eq!(tracker.get_active_ip_count("test_user").await, 0);
}

#[tokio::test]
async fn test_ip_removal() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 2).await;

    let ip1 = test_ipv4(192, 168, 1, 1);
    let ip2 = test_ipv4(192, 168, 1, 2);
    let ip3 = test_ipv4(192, 168, 1, 3);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip3).await.is_err());

    tracker.remove_ip("test_user", ip1).await;

    assert!(tracker.check_and_add("test_user", ip3).await.is_ok());
    assert_eq!(tracker.get_active_ip_count("test_user").await, 2);
}

#[tokio::test]
async fn test_no_limit() {
    let tracker = UserIpTracker::new();

    let ip1 = test_ipv4(192, 168, 1, 1);
    let ip2 = test_ipv4(192, 168, 1, 2);
    let ip3 = test_ipv4(192, 168, 1, 3);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip3).await.is_ok());

    assert_eq!(tracker.get_active_ip_count("test_user").await, 3);
}

#[tokio::test]
async fn test_multiple_users() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("user1", 2).await;
    tracker.set_user_limit("user2", 1).await;

    let ip1 = test_ipv4(192, 168, 1, 1);
    let ip2 = test_ipv4(192, 168, 1, 2);

    assert!(tracker.check_and_add("user1", ip1).await.is_ok());
    assert!(tracker.check_and_add("user1", ip2).await.is_ok());

    assert!(tracker.check_and_add("user2", ip1).await.is_ok());
    assert!(tracker.check_and_add("user2", ip2).await.is_err());
}

#[tokio::test]
async fn test_ipv6_support() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 2).await;

    let ipv4 = test_ipv4(192, 168, 1, 1);
    let ipv6 = test_ipv6();

    assert!(tracker.check_and_add("test_user", ipv4).await.is_ok());
    assert!(tracker.check_and_add("test_user", ipv6).await.is_ok());

    assert_eq!(tracker.get_active_ip_count("test_user").await, 2);
}

#[tokio::test]
async fn test_get_active_ips() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 3).await;

    let ip1 = test_ipv4(192, 168, 1, 1);
    let ip2 = test_ipv4(192, 168, 1, 2);

    tracker.check_and_add("test_user", ip1).await.unwrap();
    tracker.check_and_add("test_user", ip2).await.unwrap();

    let active_ips = tracker.get_active_ips("test_user").await;
    assert_eq!(active_ips.len(), 2);
    assert!(active_ips.contains(&ip1));
    assert!(active_ips.contains(&ip2));
}

#[tokio::test]
async fn test_stats() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("user1", 3).await;
    tracker.set_user_limit("user2", 2).await;

    let ip1 = test_ipv4(192, 168, 1, 1);
    let ip2 = test_ipv4(192, 168, 1, 2);

    tracker.check_and_add("user1", ip1).await.unwrap();
    tracker.check_and_add("user2", ip2).await.unwrap();

    let stats = tracker.get_stats().await;
    assert_eq!(stats.len(), 2);

    assert!(stats.iter().any(|(name, _, _)| name == "user1"));
    assert!(stats.iter().any(|(name, _, _)| name == "user2"));
}

#[tokio::test]
async fn test_clear_user_ips() {
    let tracker = UserIpTracker::new();
    let ip1 = test_ipv4(192, 168, 1, 1);

    tracker.check_and_add("test_user", ip1).await.unwrap();
    assert_eq!(tracker.get_active_ip_count("test_user").await, 1);

    tracker.clear_user_ips("test_user").await;
    assert_eq!(tracker.get_active_ip_count("test_user").await, 0);
}

#[tokio::test]
async fn test_is_ip_active() {
    let tracker = UserIpTracker::new();
    let ip1 = test_ipv4(192, 168, 1, 1);
    let ip2 = test_ipv4(192, 168, 1, 2);

    tracker.check_and_add("test_user", ip1).await.unwrap();

    assert!(tracker.is_ip_active("test_user", ip1).await);
    assert!(!tracker.is_ip_active("test_user", ip2).await);
}

#[tokio::test]
async fn test_load_limits_from_config() {
    let tracker = UserIpTracker::new();

    let mut config_limits = HashMap::new();
    config_limits.insert("user1".to_string(), 5);
    config_limits.insert("user2".to_string(), 3);

    tracker.load_limits(0, &config_limits).await;

    assert_eq!(tracker.get_user_limit("user1").await, Some(5));
    assert_eq!(tracker.get_user_limit("user2").await, Some(3));
    assert_eq!(tracker.get_user_limit("user3").await, None);
}

#[tokio::test]
async fn test_load_limits_replaces_previous_map() {
    let tracker = UserIpTracker::new();

    let mut first = HashMap::new();
    first.insert("user1".to_string(), 2);
    first.insert("user2".to_string(), 3);
    tracker.load_limits(0, &first).await;

    let mut second = HashMap::new();
    second.insert("user2".to_string(), 5);
    tracker.load_limits(0, &second).await;

    assert_eq!(tracker.get_user_limit("user1").await, None);
    assert_eq!(tracker.get_user_limit("user2").await, Some(5));
}

#[tokio::test]
async fn test_global_each_limit_applies_without_user_override() {
    let tracker = UserIpTracker::new();
    tracker.load_limits(2, &HashMap::new()).await;

    let ip1 = test_ipv4(172, 16, 0, 1);
    let ip2 = test_ipv4(172, 16, 0, 2);
    let ip3 = test_ipv4(172, 16, 0, 3);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip3).await.is_err());
    assert_eq!(tracker.get_user_limit("test_user").await, Some(2));
}

#[tokio::test]
async fn test_user_override_wins_over_global_each_limit() {
    let tracker = UserIpTracker::new();
    let mut limits = HashMap::new();
    limits.insert("test_user".to_string(), 1);
    tracker.load_limits(3, &limits).await;

    let ip1 = test_ipv4(172, 17, 0, 1);
    let ip2 = test_ipv4(172, 17, 0, 2);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip2).await.is_err());
    assert_eq!(tracker.get_user_limit("test_user").await, Some(1));
}

#[tokio::test]
async fn test_time_window_mode_blocks_recent_ip_churn() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 1).await;
    tracker
        .set_limit_policy(UserMaxUniqueIpsMode::TimeWindow, 30)
        .await;

    let ip1 = test_ipv4(10, 0, 0, 1);
    let ip2 = test_ipv4(10, 0, 0, 2);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    tracker.remove_ip("test_user", ip1).await;
    assert!(tracker.check_and_add("test_user", ip2).await.is_err());
}

#[tokio::test]
async fn test_combined_mode_enforces_active_and_recent_limits() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 1).await;
    tracker
        .set_limit_policy(UserMaxUniqueIpsMode::Combined, 30)
        .await;

    let ip1 = test_ipv4(10, 0, 1, 1);
    let ip2 = test_ipv4(10, 0, 1, 2);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    assert!(tracker.check_and_add("test_user", ip2).await.is_err());

    tracker.remove_ip("test_user", ip1).await;
    assert!(tracker.check_and_add("test_user", ip2).await.is_err());
}

#[tokio::test]
async fn test_time_window_expires() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 1).await;
    tracker
        .set_limit_policy(UserMaxUniqueIpsMode::TimeWindow, 1)
        .await;

    let ip1 = test_ipv4(10, 1, 0, 1);
    let ip2 = test_ipv4(10, 1, 0, 2);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    tracker.remove_ip("test_user", ip1).await;
    assert!(tracker.check_and_add("test_user", ip2).await.is_err());

    tokio::time::sleep(Duration::from_millis(1100)).await;
    assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
}

#[tokio::test]
async fn test_memory_stats_reports_queue_and_entry_counts() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 4).await;
    let ip1 = test_ipv4(10, 2, 0, 1);
    let ip2 = test_ipv4(10, 2, 0, 2);

    tracker.check_and_add("test_user", ip1).await.unwrap();
    tracker.check_and_add("test_user", ip2).await.unwrap();
    tracker.enqueue_cleanup("test_user".to_string(), ip1);

    let snapshot = tracker.memory_stats().await;
    assert_eq!(snapshot.active_users, 1);
    assert_eq!(snapshot.recent_users, 1);
    assert_eq!(snapshot.active_entries, 2);
    assert_eq!(snapshot.recent_entries, 2);
    assert_eq!(snapshot.cleanup_queue_len, 1);
}

#[tokio::test]
async fn test_compact_prunes_stale_recent_entries() {
    let tracker = UserIpTracker::new();
    tracker
        .set_limit_policy(UserMaxUniqueIpsMode::TimeWindow, 1)
        .await;

    let stale_user = "stale-user".to_string();
    let stale_ip = test_ipv4(10, 3, 0, 1);
    {
        let shard_idx = UserIpTracker::shard_idx(&stale_user);
        let mut shard = tracker.shards[shard_idx].write().await;
        shard
            .recent_ips
            .entry(stale_user.clone())
            .or_insert_with(HashMap::new)
            .insert(stale_ip, Instant::now() - Duration::from_secs(5));
    }

    tracker.last_compact_epoch_secs.store(0, Ordering::Relaxed);
    tracker
        .check_and_add("trigger-user", test_ipv4(10, 3, 0, 2))
        .await
        .unwrap();

    let shard_idx = UserIpTracker::shard_idx(&stale_user);
    let shard = tracker.shards[shard_idx].read().await;
    let stale_exists = shard
        .recent_ips
        .get(&stale_user)
        .map(|ips| ips.contains_key(&stale_ip))
        .unwrap_or(false);
    assert!(!stale_exists);
}

#[tokio::test]
async fn test_time_window_allows_same_ip_reconnect() {
    let tracker = UserIpTracker::new();
    tracker.set_user_limit("test_user", 1).await;
    tracker
        .set_limit_policy(UserMaxUniqueIpsMode::TimeWindow, 1)
        .await;

    let ip1 = test_ipv4(10, 4, 0, 1);

    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
    tracker.remove_ip("test_user", ip1).await;
    assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
}