tidaldb/tidal-server/tests/middleware.rs

//! Integration tests for the tidal-server middleware stack.
//!
//! Each test spins up an in-process router using `tower::ServiceExt::oneshot`
//! — no TCP port binding required — and verifies middleware behaviors:
//! body limits, authentication, and request-ID propagation.

use std::sync::Arc;

use axum::body::Body;
use axum::http::{Method, Request, StatusCode};
use tidal_server::{router::build_router, state::ServerState};
use tidaldb::TidalDb;
use tower::ServiceExt;

// ── Test helpers ─────────────────────────────────────────────────────────────

fn make_state() -> Arc<ServerState> {
    // Use the default schema so ranking profiles are available — same setup as
    // `run_standalone` in main.rs.
    let (schema, profiles) = tidal_server::config::load_schema(None).unwrap();
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema)
        .with_profiles(profiles)
        .open()
        .unwrap();
    Arc::new(ServerState::new(db))
}

fn make_app(api_key: Option<&str>) -> axum::Router {
    let key = api_key.map(|k| Arc::from(k));
    build_router(make_state(), key)
}

// ── Auth tests ────────────────────────────────────────────────────────────────

/// /health is public — no Bearer token required.
#[tokio::test]
async fn health_no_auth_required() {
    let app = make_app(Some("secret"));
    let response = app
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

/// Protected routes (e.g. /feed) require a Bearer token when a key is set.
#[tokio::test]
async fn protected_route_requires_auth() {
    let app = make_app(Some("secret"));
    let response = app
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/feed")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}

/// A valid Bearer token grants access to protected routes.
#[tokio::test]
async fn protected_route_accepts_valid_key() {
    let app = make_app(Some("secret"));
    let response = app
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/feed")
                .header("Authorization", "Bearer secret")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

/// RFC 7235: auth scheme is case-insensitive — `bearer` should work like `Bearer`.
#[tokio::test]
async fn lowercase_bearer_accepted() {
    let app = make_app(Some("secret"));
    let response = app
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/feed")
                .header("Authorization", "bearer secret")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

/// An incorrect Bearer token is rejected with 401.
#[tokio::test]
async fn wrong_key_rejected() {
    let app = make_app(Some("secret"));
    let response = app
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/feed")
                .header("Authorization", "Bearer wrong")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}

/// When no API key is configured, all routes are accessible without auth.
#[tokio::test]
async fn no_auth_mode_allows_all_routes() {
    let app = make_app(None);
    let response = app
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/feed")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}

// ── Middleware behavior tests ─────────────────────────────────────────────────

/// Bodies exceeding 2MB on protected routes are rejected with 413.
#[tokio::test]
async fn body_too_large_returns_413() {
    let app = make_app(Some("secret"));
    // 3 MB exceeds the 2 MB BODY_LIMIT_BYTES constant in router.rs.
    let big_body = vec![0u8; 3 * 1024 * 1024];
    let response = app
        .oneshot(
            Request::builder()
                .method(Method::POST)
                .uri("/embeddings")
                .header("Authorization", "Bearer secret")
                .header("Content-Type", "application/json")
                .body(Body::from(big_body))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
}

/// x-request-id is present in every response and increments per request.
#[tokio::test]
async fn x_request_id_increments() {
    // Both calls share the same App (and therefore the same AtomicU64 counter).
    let app = make_app(None);
    let r1 = app
        .clone()
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    let r2 = app
        .oneshot(
            Request::builder()
                .method(Method::GET)
                .uri("/health")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    let id1: u64 = r1
        .headers()
        .get("x-request-id")
        .expect("x-request-id missing from first response")
        .to_str()
        .unwrap()
        .parse()
        .unwrap();
    let id2: u64 = r2
        .headers()
        .get("x-request-id")
        .expect("x-request-id missing from second response")
        .to_str()
        .unwrap()
        .parse()
        .unwrap();

    assert!(id2 > id1, "x-request-id should increase: {id1} then {id2}");
}