athena_rs 2.0.2

use actix_web::dev::ServiceResponse;
use actix_web::web::Data;
use actix_web::{App, http::StatusCode, test};
use athena_rs::AppState;
use athena_rs::api::admin;
use athena_rs::api::auth::{api_key_query_param_used, extract_api_key};
use athena_rs::api::gateway::fetch::fetch_data_route;
use athena_rs::drivers::postgresql::sqlx_driver::PostgresClientRegistry;
use moka::future::Cache;
use reqwest::Client;
use serde_json::{Value, json};
use std::sync::Arc;
use std::time::Instant;

fn build_test_state(auth_client_name: Option<&str>) -> Data<AppState> {
    let cache: Arc<Cache<String, Value>> = Arc::new(Cache::builder().build());
    let immortal_cache: Arc<Cache<String, Value>> = Arc::new(Cache::builder().build());
    let client: Client = Client::builder()
        .build()
        .expect("Failed to build HTTP client");
    let pg_registry: Arc<PostgresClientRegistry> = Arc::new(PostgresClientRegistry::empty());
    let jdbc_pool_cache: Arc<Cache<String, sqlx::Pool<sqlx::Postgres>>> =
        Arc::new(Cache::builder().max_capacity(64).build());

    Data::new(AppState {
        cache,
        immortal_cache,
        client,
        process_start_time_seconds: 0,
        process_started_at: Instant::now(),
        pg_registry,
        jdbc_pool_cache,
        gateway_force_camel_case_to_snake_case: false,
        gateway_auto_cast_uuid_filter_values_to_text: true,
        pipeline_registry: None,
        logging_client_name: None,
        gateway_auth_client_name: auth_client_name.map(str::to_string),
        gateway_jdbc_allow_private_hosts: true,
        gateway_jdbc_allowed_hosts: vec![],
        prometheus_metrics_enabled: false,
        metrics_state: Arc::new(athena_rs::api::metrics::MetricsState::new()),
    })
}

#[actix_web::test]
async fn admin_routes_require_static_admin_key() {
    let state = build_test_state(None);
    let app = test::init_service(
        App::new()
            .app_data(state.clone())
            .configure(admin::services),
    )
    .await;

    let req = test::TestRequest::get().uri("/admin/api-keys").to_request();
    let resp: ServiceResponse = test::call_service(&app, req).await;

    assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}

#[actix_web::test]
async fn admin_routes_fail_when_auth_store_is_not_configured() {
    unsafe {
        std::env::set_var("ATHENA_KEY_12", "test-admin-key");
    }
    let state = build_test_state(None);
    let app = test::init_service(
        App::new()
            .app_data(state.clone())
            .configure(admin::services),
    )
    .await;

    let req = test::TestRequest::get()
        .uri("/admin/api-keys")
        .insert_header(("X-Athena-Key", "test-admin-key"))
        .to_request();
    let resp: ServiceResponse = test::call_service(&app, req).await;

    assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[actix_web::test]
async fn gateway_fetch_fails_closed_when_auth_store_is_configured_but_unavailable() {
    let state = build_test_state(Some("auth_store"));
    let app =
        test::init_service(App::new().app_data(state.clone()).service(fetch_data_route)).await;

    let req = test::TestRequest::post()
        .uri("/gateway/data")
        .insert_header(("X-Athena-Client", "reporting"))
        .set_json(json!({
            "table_name": "users"
        }))
        .to_request();
    let resp: ServiceResponse = test::call_service(&app, req).await;

    assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
    let body: Value = test::read_body_json(resp).await;
    assert_eq!(body["status"], "error");
}

#[actix_web::test]
async fn extract_api_key_reads_bearer_before_query_param() {
    let req = test::TestRequest::with_uri("/gateway/data?api_key=query-key")
        .insert_header(("Authorization", "Bearer header-key"))
        .to_http_request();

    assert_eq!(extract_api_key(&req).as_deref(), Some("header-key"));
}

#[actix_web::test]
async fn extract_api_key_supports_legacy_apikey_query_param() {
    let req = test::TestRequest::with_uri("/gateway/data?apikey=legacy-key").to_http_request();

    assert_eq!(extract_api_key(&req).as_deref(), Some("legacy-key"));
    assert!(api_key_query_param_used(&req));
}

#[actix_web::test]
async fn api_key_query_param_used_is_false_for_header_auth() {
    let req = test::TestRequest::default()
        .insert_header(("x-api-key", "header-key"))
        .to_http_request();

    assert!(!api_key_query_param_used(&req));
}