union_square 0.2.0

A proxy/wire-tap service for making LLM calls and recording everything that happens in a session for later analysis and test-case extraction
Documentation
//! Tests for proxy middleware layers

#[cfg(test)]
mod tests {
    use axum::body::Body;
    use axum::http::header;
    use http::Request;
    use uuid::Uuid;

    mod request_id_middleware {
        use super::*;

        #[tokio::test]
        async fn test_request_id_generation() {
            // Test that request ID middleware generates a valid UUID v7
            let _request = Request::builder()
                .method("GET")
                .uri("/test")
                .body(Body::empty())
                .unwrap();

            // TODO: Apply request ID middleware
            // let response = middleware.oneshot(request).await.unwrap();

            // Should have X-Request-ID header with valid UUID v7
            // assert!(response.headers().contains_key("x-request-id"));
            // let request_id = response.headers().get("x-request-id").unwrap();
            // let uuid = Uuid::parse_str(request_id.to_str().unwrap()).unwrap();
            // assert_eq!(uuid.get_version_num(), 7);
        }

        #[tokio::test]
        async fn test_request_id_passthrough() {
            // Test that existing request IDs are preserved
            let existing_id = Uuid::now_v7().to_string();
            let _request = Request::builder()
                .method("GET")
                .uri("/test")
                .header("x-request-id", &existing_id)
                .body(Body::empty())
                .unwrap();

            // TODO: Apply request ID middleware
            // let response = middleware.oneshot(request).await.unwrap();

            // Should preserve the existing request ID
            // assert_eq!(
            //     response.headers().get("x-request-id").unwrap().to_str().unwrap(),
            //     existing_id
            // );
        }

        #[tokio::test]
        async fn test_request_id_propagation() {
            // Test that request ID is propagated through the request chain
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .body(Body::from("test body"))
                .unwrap();

            // TODO: Apply middleware stack with request ID
            // let response = middleware_stack.oneshot(request).await.unwrap();

            // Request ID should be available in both request and response
            // assert!(response.headers().contains_key("x-request-id"));
        }
    }

    mod auth_middleware {
        use super::*;

        #[tokio::test]
        async fn test_valid_api_key() {
            // Test that valid API keys are accepted
            let valid_key = "valid-api-key-123";
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header(header::AUTHORIZATION, format!("Bearer {valid_key}"))
                .body(Body::empty())
                .unwrap();

            // TODO: Apply auth middleware with configured valid keys
            // let response = auth_middleware.oneshot(request).await.unwrap();

            // Should pass through with 200 OK
            // assert_eq!(response.status(), StatusCode::OK);
        }

        #[tokio::test]
        async fn test_missing_api_key() {
            // Test that missing API keys are rejected
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .body(Body::empty())
                .unwrap();

            // TODO: Apply auth middleware
            // let response = auth_middleware.oneshot(request).await.unwrap();

            // Should return 401 Unauthorized
            // assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
        }

        #[tokio::test]
        async fn test_invalid_api_key() {
            // Test that invalid API keys are rejected
            let invalid_key = "invalid-api-key";
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header(header::AUTHORIZATION, format!("Bearer {invalid_key}"))
                .body(Body::empty())
                .unwrap();

            // TODO: Apply auth middleware
            // let response = auth_middleware.oneshot(request).await.unwrap();

            // Should return 401 Unauthorized
            // assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
        }

        #[tokio::test]
        async fn test_malformed_auth_header() {
            // Test that malformed auth headers are rejected
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header(header::AUTHORIZATION, "NotBearer token")
                .body(Body::empty())
                .unwrap();

            // TODO: Apply auth middleware
            // let response = auth_middleware.oneshot(request).await.unwrap();

            // Should return 401 Unauthorized
            // assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
        }

        #[tokio::test]
        async fn test_auth_bypass_for_health_check() {
            // Test that health check endpoint bypasses auth
            let _request = Request::builder()
                .method("GET")
                .uri("/health")
                .body(Body::empty())
                .unwrap();

            // TODO: Apply auth middleware
            // let response = auth_middleware.oneshot(request).await.unwrap();

            // Should pass through without auth
            // assert_eq!(response.status(), StatusCode::OK);
        }
    }

    mod error_handling_middleware {
        use super::*;

        #[tokio::test]
        async fn test_proxy_error_formatting() {
            // Test that ProxyError is properly formatted in responses
            // This would test the error handling middleware converting
            // internal errors to proper HTTP responses

            // TODO: Create a service that returns a ProxyError
            // let failing_service = tower::service_fn(|_| async {
            //     Err::<Response<Body>, ProxyError>(ProxyError::RequestTimeout(Duration::from_secs(30)))
            // });

            // TODO: Wrap with error handling middleware
            // let response = error_middleware.oneshot(request).await.unwrap();

            // Should return appropriate status and error message
            // assert_eq!(response.status(), StatusCode::REQUEST_TIMEOUT);
            // let body = hyper::body::to_bytes(response.into_body()).await.unwrap();
            // let error_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
            // assert_eq!(error_json["error"]["type"], "request_timeout");
        }

        #[tokio::test]
        async fn test_panic_recovery() {
            // Test that panics are caught and converted to 500 errors
            // TODO: Create a service that panics
            // let panicking_service = tower::service_fn(|_| async {
            //     panic!("Unexpected error!");
            // });

            // TODO: Wrap with panic recovery middleware
            // let response = panic_middleware.oneshot(request).await.unwrap();

            // Should return 500 Internal Server Error
            // assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
        }

        #[tokio::test]
        async fn test_error_correlation() {
            // Test that errors include request ID for correlation
            let request_id = Uuid::now_v7().to_string();
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header("x-request-id", &request_id)
                .body(Body::empty())
                .unwrap();

            // TODO: Apply error handling with request ID correlation
            // let response = error_middleware.oneshot(request).await.unwrap();

            // Error response should include request ID
            // let body = hyper::body::to_bytes(response.into_body()).await.unwrap();
            // let error_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
            // assert_eq!(error_json["request_id"], request_id);
        }
    }

    mod combined_middleware {
        use super::*;

        #[tokio::test]
        async fn test_middleware_ordering() {
            // Test that middleware layers are applied in correct order:
            // 1. Request ID (first, so all logs have request ID)
            // 2. Error handling (catches errors from auth and below)
            // 3. Auth (validates before processing)
            // 4. Actual proxy handler

            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header(header::AUTHORIZATION, "Bearer valid-key")
                .body(Body::from("test"))
                .unwrap();

            // TODO: Apply full middleware stack
            // let response = middleware_stack.oneshot(request).await.unwrap();

            // Should have request ID from first middleware
            // assert!(response.headers().contains_key("x-request-id"));

            // Should pass auth and return success
            // assert_eq!(response.status(), StatusCode::OK);
        }

        #[tokio::test]
        async fn test_middleware_error_propagation() {
            // Test that errors from inner middleware are properly handled
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                // Missing auth header
                .body(Body::empty())
                .unwrap();

            // TODO: Apply full middleware stack
            // let response = middleware_stack.oneshot(request).await.unwrap();

            // Should have request ID even on auth failure
            // assert!(response.headers().contains_key("x-request-id"));

            // Should return 401 from auth middleware
            // assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
        }
    }

    mod logging_middleware {
        use super::*;

        #[tokio::test]
        async fn test_request_logging() {
            // Test that requests are logged with appropriate details
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header("x-request-id", Uuid::now_v7().to_string())
                .header(header::CONTENT_LENGTH, "100")
                .body(Body::from("test body"))
                .unwrap();

            // TODO: Apply logging middleware
            // Should log: method, path, request_id, content_length
            // let response = logging_middleware.oneshot(request).await.unwrap();

            // Verify log output contains expected fields
            // (Would need to capture logs in test)
        }

        #[tokio::test]
        async fn test_response_logging() {
            // Test that responses are logged with timing info
            let _request = Request::builder()
                .method("GET")
                .uri("/health")
                .body(Body::empty())
                .unwrap();

            // TODO: Apply logging middleware with timing
            // let start = Instant::now();
            // let response = logging_middleware.oneshot(request).await.unwrap();
            // let duration = start.elapsed();

            // Should log: status, duration_ms, request_id
            // assert!(duration.as_millis() > 0);
        }
    }

    mod rate_limiting_middleware {
        use super::*;

        #[tokio::test]
        async fn test_rate_limit_per_api_key() {
            // Test that rate limiting is applied per API key
            let api_key = "test-key";

            // TODO: Create rate limiter with low limit for testing
            // let rate_limiter = RateLimiter::new(2, Duration::from_secs(1));

            // First two requests should succeed
            for _ in 0..2 {
                let _request = Request::builder()
                    .method("POST")
                    .uri("/api/v1/completion")
                    .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
                    .body(Body::empty())
                    .unwrap();

                // let response = rate_limit_middleware.oneshot(request).await.unwrap();
                // assert_eq!(response.status(), StatusCode::OK);
            }

            // Third request should be rate limited
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
                .body(Body::empty())
                .unwrap();

            // let response = rate_limit_middleware.oneshot(request).await.unwrap();
            // assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
            // assert!(response.headers().contains_key("retry-after"));
        }

        #[tokio::test]
        async fn test_rate_limit_headers() {
            // Test that rate limit headers are included
            let _request = Request::builder()
                .method("POST")
                .uri("/api/v1/completion")
                .header(header::AUTHORIZATION, "Bearer test-key")
                .body(Body::empty())
                .unwrap();

            // TODO: Apply rate limiting middleware
            // let response = rate_limit_middleware.oneshot(request).await.unwrap();

            // Should include rate limit headers
            // assert!(response.headers().contains_key("x-ratelimit-limit"));
            // assert!(response.headers().contains_key("x-ratelimit-remaining"));
            // assert!(response.headers().contains_key("x-ratelimit-reset"));
        }
    }
}