athena_rs 2.9.1

use actix_web::http::header::CACHE_CONTROL;
use actix_web::test::TestRequest;
use actix_web::web::Data;
use athena_rs::api::cache::{
    cache_control::is_cache_control_no_cache,
    check::{
        check_cache_control_and_get_response, check_cache_control_and_get_response_v2,
        get_cached_response,
    },
    hydrate::hydrate_cache_and_return_json,
    rehydrate::check_rehydrate,
};
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;

use athena_rs::AppState;

fn build_test_app_state() -> Data<AppState> {
    let cache: Cache<String, Value> = Cache::builder().build();
    let immortal: Cache<String, Value> = Cache::builder().build();
    let jdbc_pool_cache = Arc::new(Cache::builder().max_capacity(64).build());
    Data::new(AppState {
        cache: Arc::new(cache),
        immortal_cache: Arc::new(immortal),
        client: Client::new(),
        process_start_time_seconds: 0,
        process_started_at: Instant::now(),
        pg_registry: Arc::new(PostgresClientRegistry::empty()),
        jdbc_pool_cache,
        gateway_force_camel_case_to_snake_case: false,
        gateway_auto_cast_uuid_filter_values_to_text: true,
        gateway_allow_schema_names_prefixed_as_table_name: true,
        pipeline_registry: None,
        logging_client_name: None,
        gateway_auth_client_name: None,
        gateway_jdbc_allow_private_hosts: true,
        gateway_jdbc_allowed_hosts: vec![],
        gateway_resilience_timeout_secs: 30,
        gateway_resilience_read_max_retries: 1,
        gateway_resilience_initial_backoff_ms: 100,
        prometheus_metrics_enabled: false,
        metrics_state: Arc::new(athena_rs::api::metrics::MetricsState::new()),
    })
}

#[actix_web::test]
async fn cache_control_no_cache_true_when_header_present() {
    let req = TestRequest::default()
        .insert_header(("Cache-Control", "no-cache"))
        .to_http_request();
    assert!(is_cache_control_no_cache(&req));
}

#[actix_web::test]
async fn cache_control_no_cache_false_when_header_missing() {
    let req = TestRequest::default().to_http_request();
    assert!(!is_cache_control_no_cache(&req));
}

#[actix_web::test]
async fn cache_control_no_cache_false_other_directives() {
    let req = TestRequest::default()
        .insert_header(("Cache-Control", "max-age=60"))
        .to_http_request();
    assert!(!is_cache_control_no_cache(&req));
}

#[actix_web::test]
async fn rehydrate_returns_true_when_present() {
    let req = TestRequest::default()
        .uri("/path?rehydrate=true")
        .to_http_request();
    assert!(check_rehydrate(&req).await);
}

#[actix_web::test]
async fn rehydrate_returns_false_when_missing() {
    let req = TestRequest::default()
        .uri("/path?foo=bar")
        .to_http_request();
    assert!(!check_rehydrate(&req).await);
}

#[tokio::test]
async fn hydrate_inserts_and_returns() {
    let data = build_test_app_state();
    let key = "test_key".to_string();
    let payload = vec![json!({"a": 1}), json!({"b": 2})];

    let returned = hydrate_cache_and_return_json(data.clone(), key.clone(), payload.clone()).await;
    assert_eq!(returned, payload);
    let cached = data.cache.get(&key).await;
    assert!(cached.is_some());
    assert_eq!(cached.unwrap(), Value::Array(payload));
}

#[tokio::test]
async fn hydrate_overwrites_existing_key() {
    let data = build_test_app_state();
    let key = "key2".to_string();
    hydrate_cache_and_return_json(data.clone(), key.clone(), vec![json!({"v": 1})]).await;
    hydrate_cache_and_return_json(data.clone(), key.clone(), vec![json!({"v": 2})]).await;
    let cached = data.cache.get(&key).await.unwrap();
    assert_eq!(cached, Value::Array(vec![json!({"v": 2})]));
}

#[tokio::test]
async fn check_get_cached_response_returns_value() {
    let data = build_test_app_state();
    let key = "k".to_string();
    data.cache.insert(key.clone(), json!({"x": 1})).await;
    let maybe = get_cached_response(data, &key).await;
    assert!(maybe.is_some());
}

#[tokio::test]
async fn check_cache_control_no_cache_skips() {
    let data = build_test_app_state();
    let key = "k2".to_string();
    data.cache.insert(key.clone(), json!({"y": 2})).await;

    let req = TestRequest::default()
        .insert_header((CACHE_CONTROL, "no-cache"))
        .to_http_request();

    let got = check_cache_control_and_get_response(&req, data, &key).await;
    assert!(got.is_none());
}

#[tokio::test]
async fn check_v2_flattens_single_item_array() {
    let data = build_test_app_state();
    let key = "k3".to_string();
    data.cache.insert(key.clone(), json!([{"z": 3}])).await;

    let req = TestRequest::default().to_http_request();
    let resp = check_cache_control_and_get_response_v2(&req, data, &key).await;
    assert!(resp.is_some());
}

#[tokio::test]
async fn check_v2_cache_miss_returns_none() {
    let data = build_test_app_state();
    let key = "missing".to_string();

    let req = TestRequest::default().to_http_request();
    let resp = check_cache_control_and_get_response_v2(&req, data, &key).await;
    assert!(resp.is_none());
}

#[tokio::test]
async fn check_strips_nulls_when_header_set() {
    let data = build_test_app_state();
    let key = "strip".to_string();
    data.cache
        .insert(key.clone(), json!({"data": [{"a": 1, "b": null}]}))
        .await;
    let req = TestRequest::default()
        .insert_header(("X-Strip-Nulls", "true"))
        .to_http_request();
    let resp = check_cache_control_and_get_response(&req, data, &key).await;
    assert!(resp.is_some());
}