imphnen-backend-service/tests/src/rate_limiting_middleware_test.rs

#[cfg(test)]
mod rate_limiting_middleware_tests {
    use axum::{http::Request, middleware::Next, response::Response};
    use imphnen_libs::{AppState, environment::Environment};
    use imphnen_middleware::rate_limiting_middleware::{
        RateLimitConfig, RateLimitStore, TokenBucket, create_rate_limiting_middleware,
        auth_rate_limiting_middleware,
    };
    use std::{sync::Arc, time::Duration};
    use tower::ServiceExt;

    #[tokio::test]
    async fn test_token_bucket_basic_functionality() {
        let bucket = TokenBucket::new(5, 2); // Capacity 5, refill 2 per second
        
        // Should have full tokens initially
        assert_eq!(bucket.tokens, 5);
        
        // Consume some tokens
        assert!(bucket.try_consume());
        assert_eq!(bucket.tokens, 4);
        
        assert!(bucket.try_consume());
        assert_eq!(bucket.tokens, 3);
        
        assert!(bucket.try_consume());
        assert_eq!(bucket.tokens, 2);
        
        assert!(bucket.try_consume());
        assert_eq!(bucket.tokens, 1);
        
        assert!(bucket.try_consume());
        assert_eq!(bucket.tokens, 0);
        
        // Should not consume when empty
        assert!(!bucket.try_consume());
        assert_eq!(bucket.tokens, 0);
    }

    #[tokio::test]
    async fn test_token_bucket_refill() {
        let mut bucket = TokenBucket::new(3, 1); // Capacity 3, refill 1 per second
        
        // Consume all tokens
        for _ in 0..3 {
            assert!(bucket.try_consume());
        }
        
        assert!(!bucket.try_consume());
        assert_eq!(bucket.tokens, 0);
        
        // Wait for 1 second to allow refill
        tokio::time::sleep(Duration::from_secs(1)).await;
        
        // Should have 1 token after refill
        bucket.refill_tokens();
        assert_eq!(bucket.tokens, 1);
        
        // Consume the refilled token
        assert!(bucket.try_consume());
        assert_eq!(bucket.tokens, 0);
        
        // Wait another second
        tokio::time::sleep(Duration::from_secs(1)).await;
        
        // Should have another token
        bucket.refill_tokens();
        assert_eq!(bucket.tokens, 1);
    }

    #[tokio::test]
    async fn test_rate_limit_store_basic() {
        let config = RateLimitConfig::test();
        let store = Arc::new(RateLimitStore::new(config));
        
        let client_ip = "127.0.0.1";
        
        // First request should succeed
        let result = store.check_limit(client_ip).await;
        assert!(result.is_ok());
        
        // Multiple requests should succeed within limits
        for _ in 0..config.bucket_size {
            let result = store.check_limit(client_ip).await;
            assert!(result.is_ok());
        }
        
        // Next request should fail
        let result = store.check_limit(client_ip).await;
        assert!(result.is_err());
        assert_eq!(result.unwrap_err(), axum::http::StatusCode::TOO_MANY_REQUESTS);
    }

    #[tokio::test]
    async fn test_rate_limit_store_window_reset() {
        let config = RateLimitConfig {
            max_requests: 10,
            window_duration: Duration::from_secs(2),
            bucket_size: 2,
            refill_rate: 1,
        };
        let store = Arc::new(RateLimitStore::new(config));
        
        let client_ip = "127.0.0.1";
        
        // Consume all tokens
        assert!(store.check_limit(client_ip).await.is_ok());
        assert!(store.check_limit(client_ip).await.is_ok());
        assert!(store.check_limit(client_ip).await.is_err());
        
        // Wait for window to reset
        tokio::time::sleep(Duration::from_secs(3)).await;
        
        // Should be able to make requests again
        assert!(store.check_limit(client_ip).await.is_ok());
        assert!(store.check_limit(client_ip).await.is_ok());
        assert!(store.check_limit(client_ip).await.is_err());
    }

    #[tokio::test]
    async fn test_different_clients_have_separate_limits() {
        let config = RateLimitConfig::test();
        let store = Arc::new(RateLimitStore::new(config));
        
        let client_ip_1 = "127.0.0.1";
        let client_ip_2 = "127.0.0.2";
        
        // Client 1 should be able to make requests
        for _ in 0..config.bucket_size {
            assert!(store.check_limit(client_ip_1).await.is_ok());
        }
        assert!(store.check_limit(client_ip_1).await.is_err());
        
        // Client 2 should still be able to make requests
        for _ in 0..config.bucket_size {
            assert!(store.check_limit(client_ip_2).await.is_ok());
        }
        assert!(store.check_limit(client_ip_2).await.is_err());
    }

    #[tokio::test]
    async fn test_auth_rate_limiting_middleware_success() {
        // Create a mock AppState with test environment
        let state = AppState {
            surrealdb_ws: Default::default(),
            surrealdb_mem: Default::default(),
            user_lookup_service: Default::default(),
            auth_repository: Default::default(),
            env: Environment::Test,
        };

        // Create a mock request to /auth/login
        let mut request = Request::builder()
            .uri("/v1/auth/login")
            .header("x-forwarded-for", "127.0.0.1")
            .body(())
            .unwrap();

        // Create a mock next service
        let next = Next::new(|req| async move {
            let response = Response::builder()
                .status(200)
                .body("Login successful")
                .unwrap();
            Ok::<_, axum::http::StatusCode>((req, response))
        });

        // Call the middleware
        let result = auth_rate_limiting_middleware(
            axum::Extension(state.clone()),
            request,
            next,
        ).await;

        // Should succeed
        assert!(result.is_ok());
        let response = result.unwrap();
        assert_eq!(response.status(), 200);
    }

    #[tokio::test]
    async fn test_auth_rate_limiting_middleware_429() {
        // Create test configuration with very low limits for testing
        let config = RateLimitConfig {
            max_requests: 1,
            window_duration: Duration::from_secs(10),
            bucket_size: 1,
            refill_rate: 1,
        };

        // Create a mock AppState with test environment
        let state = AppState {
            surrealdb_ws: Default::default(),
            surrealdb_mem: Default::default(),
            user_lookup_service: Default::default(),
            auth_repository: Default::default(),
            env: Environment::Test,
        };

        // Create a mock request to /auth/login
        let mut request = Request::builder()
            .uri("/v1/auth/login")
            .header("x-forwarded-for", "127.0.0.1")
            .body(())
            .unwrap();

        // Create a mock next service
        let next = Next::new(|req| async move {
            let response = Response::builder()
                .status(200)
                .body("Login successful")
                .unwrap();
            Ok::<_, axum::http::StatusCode>((req, response))
        });

        // First request should succeed
        let result = auth_rate_limiting_middleware(
            axum::Extension(state.clone()),
            request.clone(),
            next.clone(),
        ).await;
        assert!(result.is_ok());

        // Second request should fail with 429
        let result = auth_rate_limiting_middleware(
            axum::Extension(state),
            request,
            next,
        ).await;
        assert!(result.is_ok());
        let response = result.unwrap();
        assert_eq!(response.status(), 429);
        assert_eq!(response.headers().get("Retry-After").unwrap(), "60");
    }

    #[tokio::test]
    async fn test_non_auth_endpoints_not_rate_limited() {
        // Create a mock AppState with test environment
        let state = AppState {
            surrealdb_ws: Default::default(),
            surrealdb_mem: Default::default(),
            user_lookup_service: Default::default(),
            auth_repository: Default::default(),
            env: Environment::Test,
        };

        // Create a mock request to a non-auth endpoint
        let mut request = Request::builder()
            .uri("/v1/users/me")
            .header("x-forwarded-for", "127.0.0.1")
            .body(())
            .unwrap();

        // Create a mock next service
        let next = Next::new(|req| async move {
            let response = Response::builder()
                .status(200)
                .body("User data")
                .unwrap();
            Ok::<_, axum::http::StatusCode>((req, response))
        });

        // Call the middleware - should not apply rate limiting
        let result = auth_rate_limiting_middleware(
            axum::Extension(state),
            request,
            next,
        ).await;

        // Should succeed
        assert!(result.is_ok());
        let response = result.unwrap();
        assert_eq!(response.status(), 200);
    }

    #[tokio::test]
    async fn test_environment_specific_configurations() {
        // Test development config
        let dev_config = RateLimitConfig::development();
        assert_eq!(dev_config.max_requests, 100);
        assert_eq!(dev_config.bucket_size, 50);
        assert_eq!(dev_config.refill_rate, 10);

        // Test production config
        let prod_config = RateLimitConfig::production();
        assert_eq!(prod_config.max_requests, 10);
        assert_eq!(prod_config.bucket_size, 5);
        assert_eq!(prod_config.refill_rate, 1);

        // Test test config
        let test_config = RateLimitConfig::test();
        assert_eq!(test_config.max_requests, 1000);
        assert_eq!(test_config.bucket_size, 100);
        assert_eq!(test_config.refill_rate, 20);
    }
}