Aether/apps/aether-gateway/src/data/tests.rs

use std::sync::Arc;
use std::time::{Duration, Instant};

use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
    InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data::repository::users::{
    InMemoryUserReadRepository, StoredUserAuthRecord, StoredUserPreferenceRecord,
};
use aether_data::repository::video_tasks::InMemoryVideoTaskRepository;
use aether_data::{DataLayerError, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
use aether_data_contracts::repository::candidate_selection::{
    StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
    RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::provider_catalog::{
    StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::repository::video_tasks::{
    UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus, VideoTaskWriteRepository,
};
use aether_scheduler_core::{
    enumerate_minimal_candidate_selection, EnumerateMinimalCandidateSelectionInput,
    SchedulerAuthConstraints,
};
use serde_json::json;

use super::{GatewayDataConfig, GatewayDataState};
use crate::AppState;

#[test]
fn disabled_gateway_data_state_has_no_backends() {
    let state = GatewayDataState::from_config(GatewayDataConfig::disabled())
        .expect("disabled config should build");

    assert!(!state.has_backends());
    assert!(!state.has_auth_api_key_reader());
    assert!(!state.has_minimal_candidate_selection_reader());
    assert!(!state.has_request_candidate_reader());
    assert!(!state.has_provider_catalog_reader());
    assert!(!state.has_proxy_node_reader());
    assert!(!state.has_proxy_node_writer());
    assert!(!state.has_usage_reader());
    assert!(!state.has_video_task_reader());
}

#[test]
fn maintenance_pool_pressure_keeps_idle_reserve_for_foreground_work() {
    let pool_can_still_grow = aether_data::DatabasePoolSummary {
        driver: DatabaseDriver::Postgres,
        checked_out: 6,
        pool_size: 6,
        idle: 0,
        max_connections: 20,
        usage_rate: 30.0,
    };
    assert!(
        !GatewayDataState::database_pool_summary_under_maintenance_pressure(&pool_can_still_grow)
    );

    let reserve_idle_left = aether_data::DatabasePoolSummary {
        driver: DatabaseDriver::Postgres,
        checked_out: 18,
        pool_size: 20,
        idle: 2,
        max_connections: 20,
        usage_rate: 90.0,
    };
    assert!(GatewayDataState::database_pool_summary_under_maintenance_pressure(&reserve_idle_left));

    let above_idle_reserve = aether_data::DatabasePoolSummary {
        driver: DatabaseDriver::Postgres,
        checked_out: 17,
        pool_size: 20,
        idle: 3,
        max_connections: 20,
        usage_rate: 85.0,
    };
    assert!(
        !GatewayDataState::database_pool_summary_under_maintenance_pressure(&above_idle_reserve)
    );

    let idle = aether_data::DatabasePoolSummary {
        driver: DatabaseDriver::Postgres,
        checked_out: 0,
        pool_size: 4,
        idle: 4,
        max_connections: 20,
        usage_rate: 0.0,
    };
    assert!(!GatewayDataState::database_pool_summary_under_maintenance_pressure(&idle));
}

#[test]
fn usage_worker_pool_pressure_only_defers_near_pool_exhaustion() {
    let comfortable = aether_data::DatabasePoolSummary {
        driver: DatabaseDriver::Postgres,
        checked_out: 56,
        pool_size: 64,
        idle: 8,
        max_connections: 64,
        usage_rate: 87.5,
    };
    assert!(!GatewayDataState::database_pool_summary_under_usage_worker_pressure(&comfortable));

    let last_idle_left = aether_data::DatabasePoolSummary {
        driver: DatabaseDriver::Postgres,
        checked_out: 63,
        pool_size: 64,
        idle: 1,
        max_connections: 64,
        usage_rate: 98.4375,
    };
    assert!(GatewayDataState::database_pool_summary_under_usage_worker_pressure(&last_idle_left));

    let exhausted = aether_data::DatabasePoolSummary {
        driver: DatabaseDriver::Postgres,
        checked_out: 64,
        pool_size: 64,
        idle: 0,
        max_connections: 64,
        usage_rate: 100.0,
    };
    assert!(GatewayDataState::database_pool_summary_under_usage_worker_pressure(&exhausted));
}

#[test]
fn maintenance_pool_pressure_deferral_has_timeout() {
    let mut deferred_since = None;
    assert!(
        GatewayDataState::should_defer_maintenance_for_pool_pressure_state(
            true,
            &mut deferred_since
        )
    );
    assert!(deferred_since.is_some());

    assert!(
        !GatewayDataState::should_defer_maintenance_for_pool_pressure_state(
            false,
            &mut deferred_since
        )
    );
    assert!(deferred_since.is_none());

    let mut stale_defer = Some(Instant::now() - Duration::from_secs(31));
    assert!(
        !GatewayDataState::should_defer_maintenance_for_pool_pressure_state(true, &mut stale_defer)
    );
    assert!(stale_defer.is_none());
}

#[tokio::test]
async fn postgres_gateway_data_state_builds_video_task_reader() {
    let state = GatewayDataState::from_config(GatewayDataConfig::from_postgres_url(
        "postgres://localhost/aether",
        false,
    ))
    .expect("postgres-backed state should build");

    assert!(state.has_backends());
    assert!(state.has_auth_api_key_reader());
    assert!(state.has_minimal_candidate_selection_reader());
    assert!(state.has_request_candidate_reader());
    assert!(state.has_provider_catalog_reader());
    assert!(state.has_proxy_node_reader());
    assert!(state.has_proxy_node_writer());
    assert!(state.has_usage_reader());
    assert!(state.has_video_task_reader());
}

#[tokio::test]
async fn data_state_find_uses_configured_read_repository() {
    let repository = Arc::new(InMemoryVideoTaskRepository::default());
    repository
        .upsert(UpsertVideoTask {
            id: "task-1".to_string(),
            short_id: Some("short-task-1".to_string()),
            request_id: "request-1".to_string(),
            user_id: Some("user-1".to_string()),
            api_key_id: Some("key-1".to_string()),
            username: Some("user".to_string()),
            api_key_name: Some("primary".to_string()),
            external_task_id: Some("ext-task-1".to_string()),
            provider_id: Some("provider-1".to_string()),
            endpoint_id: Some("endpoint-1".to_string()),
            key_id: Some("provider-key-1".to_string()),
            client_api_format: Some("openai:video".to_string()),
            provider_api_format: Some("openai:video".to_string()),
            format_converted: false,
            model: Some("sora-2".to_string()),
            prompt: Some("hello".to_string()),
            original_request_body: Some(json!({"prompt": "hello"})),
            duration_seconds: Some(4),
            resolution: Some("720p".to_string()),
            aspect_ratio: Some("16:9".to_string()),
            size: Some("1280x720".to_string()),
            status: VideoTaskStatus::Queued,
            progress_percent: 0,
            progress_message: None,
            retry_count: 0,
            poll_interval_seconds: 10,
            next_poll_at_unix_secs: Some(100),
            poll_count: 0,
            max_poll_count: 360,
            created_at_unix_ms: 100,
            submitted_at_unix_secs: Some(100),
            completed_at_unix_secs: None,
            updated_at_unix_secs: 100,
            error_code: None,
            error_message: None,
            video_url: None,
            request_metadata: None,
        })
        .await
        .expect("upsert should succeed");

    let state = GatewayDataState::with_video_task_reader_for_tests(repository);

    let task = state
        .find_video_task(VideoTaskLookupKey::Id("task-1"))
        .await
        .expect("find should succeed");

    assert_eq!(task.expect("task should exist").id, "task-1");
}

#[tokio::test]
async fn app_state_wires_gateway_data_state_from_config() {
    let state = AppState::new()
        .expect("app state should build")
        .with_data_config(GatewayDataConfig::from_postgres_url(
            "postgres://localhost/aether",
            false,
        ))
        .expect("data config should wire");

    assert!(state.data.has_backends());
    assert!(state.data.has_auth_api_key_reader());
    assert!(state.data.has_minimal_candidate_selection_reader());
    assert!(state.data.has_request_candidate_reader());
    assert!(state.data.has_provider_catalog_reader());
    assert!(state.data.has_proxy_node_reader());
    assert!(state.data.has_proxy_node_writer());
    assert!(state.data.has_usage_reader());
    assert!(state.data.has_video_task_reader());
}

#[tokio::test]
async fn data_state_checks_user_uniqueness_through_user_reader() {
    let user = StoredUserAuthRecord::new(
        "user-1".to_string(),
        Some("[email protected]".to_string()),
        true,
        "alice".to_string(),
        Some("hash".to_string()),
        "user".to_string(),
        "local".to_string(),
        None,
        None,
        None,
        true,
        false,
        None,
        None,
    )
    .expect("auth user should build");
    let admin = StoredUserAuthRecord::new(
        "admin-1".to_string(),
        Some("[email protected]".to_string()),
        true,
        "admin".to_string(),
        Some("$2b$12$4qL4tdcsFwVaDTw5Ck3xzu8GpNdre56DiNR6Dnw7t6gCXaEnqAe7G".to_string()),
        "admin".to_string(),
        "local".to_string(),
        None,
        None,
        None,
        true,
        false,
        None,
        None,
    )
    .expect("admin user should build");
    let state = GatewayDataState::with_user_reader_for_tests(Arc::new(
        InMemoryUserReadRepository::seed_auth_users(vec![user, admin]),
    ));

    assert!(state
        .is_other_user_auth_email_taken("[email protected]", "other-user")
        .await
        .expect("email uniqueness should check"));
    assert!(!state
        .is_other_user_auth_email_taken("[email protected]", "user-1")
        .await
        .expect("same user email should not be taken"));
    assert!(!state
        .is_other_user_auth_email_taken("alice", "other-user")
        .await
        .expect("email lookup should not match username"));
    assert!(state
        .is_other_user_auth_username_taken("alice", "other-user")
        .await
        .expect("username uniqueness should check"));
    assert_eq!(
        state
            .count_active_admin_users()
            .await
            .expect("active admin count should check"),
        1
    );
    assert_eq!(
        state
            .count_active_local_admin_users_with_valid_password()
            .await
            .expect("valid local admin count should check"),
        1
    );
    let preferences = StoredUserPreferenceRecord {
        user_id: "user-1".to_string(),
        avatar_url: Some("https://example.test/avatar.png".to_string()),
        bio: Some("hello".to_string()),
        default_provider_id: None,
        default_provider_name: None,
        theme: "dark".to_string(),
        language: "en-US".to_string(),
        timezone: "UTC".to_string(),
        email_notifications: false,
        usage_alerts: true,
        announcement_notifications: false,
    };
    assert_eq!(
        state
            .write_user_preferences(&preferences)
            .await
            .expect("preferences should write through repository"),
        Some(preferences.clone())
    );
    assert_eq!(
        state
            .read_user_preferences("user-1")
            .await
            .expect("preferences should read through repository"),
        Some(preferences)
    );
}

#[tokio::test]
async fn data_state_finds_active_provider_name_through_catalog_reader() {
    let active = StoredProviderCatalogProvider::new(
        "provider-1".to_string(),
        "Provider One".to_string(),
        None,
        "openai".to_string(),
    )
    .expect("provider should build");
    let inactive = StoredProviderCatalogProvider::new(
        "provider-2".to_string(),
        "Provider Two".to_string(),
        None,
        "openai".to_string(),
    )
    .expect("provider should build")
    .with_transport_fields(false, false, false, None, None, None, None, None, None);
    let state = GatewayDataState::with_provider_catalog_reader_for_tests(Arc::new(
        InMemoryProviderCatalogReadRepository::seed(vec![active, inactive], Vec::new(), Vec::new()),
    ));

    assert_eq!(
        state
            .find_active_provider_name("provider-1")
            .await
            .expect("provider lookup should succeed"),
        Some("Provider One".to_string())
    );
    assert_eq!(
        state
            .find_active_provider_name("provider-2")
            .await
            .expect("inactive provider lookup should succeed"),
        None
    );
}

fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
    StoredAuthApiKeySnapshot::new(
        user_id.to_string(),
        "alice".to_string(),
        Some("[email protected]".to_string()),
        "user".to_string(),
        "local".to_string(),
        true,
        false,
        Some(serde_json::json!(["openai"])),
        Some(serde_json::json!(["openai:chat"])),
        Some(serde_json::json!(["gpt-4.1"])),
        api_key_id.to_string(),
        Some("default".to_string()),
        true,
        false,
        false,
        Some(60),
        Some(5),
        Some(200),
        Some(serde_json::json!(["openai"])),
        Some(serde_json::json!(["openai:chat"])),
        Some(serde_json::json!(["gpt-4.1"])),
    )
    .expect("auth snapshot should build")
}

#[tokio::test]
async fn data_state_reads_auth_api_key_snapshot_from_reader() {
    let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
        Some("hash-1".to_string()),
        sample_auth_snapshot("key-1", "user-1"),
    )]));
    let state = GatewayDataState::with_auth_api_key_reader_for_tests(repository);

    let snapshot = state
        .read_auth_api_key_snapshot("user-1", "key-1", 150)
        .await
        .expect("read should succeed")
        .expect("snapshot should exist");

    assert_eq!(snapshot.user_id, "user-1");
    assert_eq!(snapshot.api_key_id, "key-1");
    assert_eq!(snapshot.username, "alice");
    assert_eq!(
        snapshot.api_key_allowed_models,
        Some(vec!["gpt-4.1".to_string()])
    );
    assert!(snapshot.currently_usable);
}

fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
    StoredProviderCatalogProvider::new(
        "provider-1".to_string(),
        "OpenAI".to_string(),
        Some("https://openai.com".to_string()),
        "custom".to_string(),
    )
    .expect("provider should build")
}

fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
    StoredProviderCatalogEndpoint::new(
        "endpoint-1".to_string(),
        "provider-1".to_string(),
        "openai:chat".to_string(),
        Some("openai".to_string()),
        Some("chat".to_string()),
        true,
    )
    .expect("endpoint should build")
}

fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
    StoredProviderCatalogKey::new(
        "provider-key-1".to_string(),
        "provider-1".to_string(),
        "prod-key".to_string(),
        "api_key".to_string(),
        Some(serde_json::json!({"cache_1h": true})),
        true,
    )
    .expect("key should build")
}

fn sample_request_usage(request_id: &str) -> StoredRequestUsageAudit {
    StoredRequestUsageAudit::new(
        "usage-1".to_string(),
        request_id.to_string(),
        Some("user-1".to_string()),
        Some("api-key-1".to_string()),
        Some("alice".to_string()),
        Some("default".to_string()),
        "OpenAI".to_string(),
        "gpt-4.1".to_string(),
        Some("gpt-4.1-mini".to_string()),
        Some("provider-1".to_string()),
        Some("endpoint-1".to_string()),
        Some("provider-key-1".to_string()),
        Some("chat".to_string()),
        Some("openai:chat".to_string()),
        Some("openai".to_string()),
        Some("chat".to_string()),
        Some("openai:chat".to_string()),
        Some("openai".to_string()),
        Some("chat".to_string()),
        true,
        false,
        120,
        40,
        160,
        0.24,
        0.36,
        Some(200),
        None,
        None,
        Some(450),
        Some(120),
        "completed".to_string(),
        "settled".to_string(),
        100,
        101,
        Some(102),
    )
    .expect("usage should build")
}

fn sample_minimal_candidate_selection_row(
    provider_id: &str,
    provider_name: &str,
    provider_priority: i32,
    key_id: &str,
    key_name: &str,
    key_internal_priority: i32,
) -> StoredMinimalCandidateSelectionRow {
    StoredMinimalCandidateSelectionRow {
        provider_id: provider_id.to_string(),
        provider_name: provider_name.to_string(),
        provider_type: "custom".to_string(),
        provider_priority,
        provider_is_active: true,
        endpoint_id: format!("endpoint-{provider_id}"),
        endpoint_api_format: "openai:chat".to_string(),
        endpoint_api_family: Some("openai".to_string()),
        endpoint_kind: Some("chat".to_string()),
        endpoint_is_active: true,
        key_id: key_id.to_string(),
        key_name: key_name.to_string(),
        key_auth_type: "api_key".to_string(),
        key_is_active: true,
        key_api_formats: Some(vec!["openai:chat".to_string()]),
        key_allowed_models: None,
        key_capabilities: Some(serde_json::json!({"cache_1h": true})),
        key_internal_priority,
        key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 3})),
        model_id: format!("model-{provider_id}"),
        global_model_id: "global-model-1".to_string(),
        global_model_name: "gpt-4.1".to_string(),
        global_model_mappings: Some(vec!["gpt-4\\.1-.*".to_string()]),
        global_model_supports_streaming: Some(true),
        model_provider_model_name: "gpt-4.1-upstream".to_string(),
        model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
            name: "gpt-4.1-canary".to_string(),
            priority: 1,
            api_formats: Some(vec!["openai:chat".to_string()]),
            endpoint_ids: None,
            operations: None,
        }]),
        model_supports_streaming: None,
        model_is_active: true,
        model_is_available: true,
    }
}

#[tokio::test]
async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
    let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
        StoredRequestCandidate::new(
            "cand-1".to_string(),
            "req-1".to_string(),
            Some("user-1".to_string()),
            Some("api-key-1".to_string()),
            Some("alice".to_string()),
            Some("default".to_string()),
            0,
            0,
            Some("provider-1".to_string()),
            Some("endpoint-1".to_string()),
            Some("provider-key-1".to_string()),
            RequestCandidateStatus::Failed,
            None,
            false,
            Some(502),
            None,
            None,
            Some(37),
            Some(1),
            None,
            Some(serde_json::json!({"cache_1h": true})),
            100_000,
            Some(101_000),
            Some(102_000),
        )
        .expect("candidate should build"),
    ]));
    let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
        vec![sample_provider_catalog_provider()],
        vec![sample_provider_catalog_endpoint()],
        vec![sample_provider_catalog_key()],
    ));
    let state = GatewayDataState::with_decision_trace_readers_for_tests(
        request_candidates,
        provider_catalog,
    );

    let trace = state
        .read_decision_trace("req-1", true)
        .await
        .expect("trace should read")
        .expect("trace should exist");

    assert_eq!(trace.request_id, "req-1");
    assert_eq!(trace.total_candidates, 1);
    assert_eq!(trace.candidates[0].provider_name.as_deref(), Some("OpenAI"));
    assert_eq!(
        trace.candidates[0].endpoint_api_format.as_deref(),
        Some("openai:chat")
    );
    assert_eq!(
        trace.candidates[0].provider_key_auth_type.as_deref(),
        Some("api_key")
    );
    assert_eq!(
        trace.candidates[0].provider_key_capabilities,
        Some(serde_json::json!({"cache_1h": true}))
    );
}

#[tokio::test]
async fn data_state_reads_request_usage_audit_from_reader() {
    let repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
        sample_request_usage("req-usage-1"),
    ]));
    let state = GatewayDataState::with_usage_reader_for_tests(repository);

    let usage = state
        .read_request_usage_audit("req-usage-1")
        .await
        .expect("read should succeed")
        .expect("usage should exist");

    assert_eq!(usage.request_id, "req-usage-1");
    assert_eq!(usage.provider_name, "OpenAI");
    assert_eq!(usage.total_tokens, 160);
    assert_eq!(usage.total_cost_usd, 0.24);
    assert!(usage.has_format_conversion);
}

#[tokio::test]
async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
    let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
        Some("hash-1".to_string()),
        sample_auth_snapshot("api-key-1", "user-1"),
    )]));
    let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
        StoredRequestCandidate::new(
            "cand-1".to_string(),
            "req-usage-1".to_string(),
            Some("user-1".to_string()),
            Some("api-key-1".to_string()),
            Some("alice".to_string()),
            Some("default".to_string()),
            0,
            0,
            Some("provider-1".to_string()),
            Some("endpoint-1".to_string()),
            Some("provider-key-1".to_string()),
            RequestCandidateStatus::Success,
            None,
            false,
            Some(200),
            None,
            None,
            Some(37),
            Some(1),
            None,
            Some(serde_json::json!({"cache_1h": true})),
            100_000,
            Some(101_000),
            Some(102_000),
        )
        .expect("candidate should build"),
    ]));
    let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
        vec![sample_provider_catalog_provider()],
        vec![sample_provider_catalog_endpoint()],
        vec![sample_provider_catalog_key()],
    ));
    let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
        sample_request_usage("req-usage-1"),
    ]));
    let state = GatewayDataState::with_request_audit_readers_for_tests(
        auth_repository,
        request_candidates,
        provider_catalog,
        usage_repository,
    );

    let bundle = state
        .read_request_audit_bundle("req-usage-1", true, 150)
        .await
        .expect("bundle should read")
        .expect("bundle should exist");

    assert_eq!(bundle.request_id, "req-usage-1");
    assert_eq!(
        bundle
            .usage
            .as_ref()
            .and_then(|usage| usage.target_model.as_deref()),
        Some("gpt-4.1-mini")
    );
    assert_eq!(
        bundle
            .decision_trace
            .as_ref()
            .and_then(|trace| trace.candidates.first())
            .and_then(|candidate| candidate.provider_name.as_deref()),
        Some("OpenAI")
    );
    assert_eq!(
        bundle
            .auth_snapshot
            .as_ref()
            .map(|snapshot| snapshot.currently_usable),
        Some(true)
    );
}

#[tokio::test]
async fn data_state_reads_decrypted_provider_transport_snapshot() {
    let encrypted_api_key =
        encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
            .expect("api key ciphertext should build");
    let encrypted_auth_config = encrypt_python_fernet_plaintext(
        DEVELOPMENT_ENCRYPTION_KEY,
        "{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}",
    )
    .expect("auth config ciphertext should build");
    let provider = sample_provider_catalog_provider().with_transport_fields(
        true,
        false,
        true,
        Some(32),
        Some(3),
        Some(serde_json::json!({"url":"http://provider-proxy"})),
        Some(20.0),
        Some(8.0),
        Some(serde_json::json!({"region":"global"})),
    );
    let endpoint = sample_provider_catalog_endpoint()
        .with_transport_fields(
            "https://api.openai.com".to_string(),
            Some(serde_json::json!([{"action":"set","key":"x-test","value":"1"}])),
            Some(serde_json::json!([{"action":"drop","path":"stream"}])),
            Some(2),
            Some("/v1/chat/completions".to_string()),
            Some(serde_json::json!({"api_version":"v1"})),
            Some(serde_json::json!({"allow":["openai:chat"]})),
            Some(serde_json::json!({"url":"http://endpoint-proxy"})),
        )
        .expect("endpoint transport should build");
    let key = sample_provider_catalog_key()
        .with_transport_fields(
            Some(serde_json::json!(["openai:chat", "openai:responses"])),
            encrypted_api_key,
            Some(encrypted_auth_config),
            Some(serde_json::json!({"openai:chat": 0.8})),
            Some(serde_json::json!({"openai:chat": 1})),
            Some(serde_json::json!(["gpt-4.1", "gpt-4.1-mini"])),
            Some(1_800_000_000),
            Some(serde_json::json!({"node_id":"proxy-node-1"})),
            Some(serde_json::json!({"transport_profile":"chrome_136"})),
        )
        .expect("key transport should build");
    let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
        vec![provider],
        vec![endpoint],
        vec![key],
    ));
    let state = GatewayDataState::with_provider_transport_reader_for_tests(
        repository,
        DEVELOPMENT_ENCRYPTION_KEY.to_string(),
    );

    let snapshot = state
        .read_provider_transport_snapshot("provider-1", "endpoint-1", "provider-key-1")
        .await
        .expect("snapshot read should succeed")
        .expect("snapshot should exist");

    assert_eq!(snapshot.provider.name, "OpenAI");
    assert_eq!(snapshot.endpoint.base_url, "https://api.openai.com");
    assert_eq!(
        snapshot.key.api_formats,
        Some(vec![
            "openai:chat".to_string(),
            "openai:responses".to_string()
        ])
    );
    assert_eq!(snapshot.key.decrypted_api_key, "sk-live-openai");
    assert_eq!(
        snapshot.key.decrypted_auth_config.as_deref(),
        Some("{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}")
    );
}

#[tokio::test]
async fn data_state_reads_minimal_candidate_selection_with_auth_filters() {
    let candidate_selection_repository =
        Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
            sample_minimal_candidate_selection_row(
                "provider-2",
                "OtherProvider",
                20,
                "key-2",
                "key-two",
                20,
            ),
            sample_minimal_candidate_selection_row(
                "provider-1",
                "OpenAI",
                10,
                "key-1",
                "key-one",
                10,
            ),
            StoredMinimalCandidateSelectionRow {
                key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 4})),
                key_allowed_models: Some(vec!["gpt-4.1-edge".to_string()]),
                ..sample_minimal_candidate_selection_row(
                    "provider-1",
                    "OpenAI",
                    10,
                    "key-3",
                    "key-three",
                    30,
                )
            },
        ]));
    let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
        Some("hash-1".to_string()),
        sample_auth_snapshot("api-key-1", "user-1"),
    )]));
    let state = GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
        candidate_selection_repository,
        auth_repository,
    );
    let auth_snapshot = state
        .read_auth_api_key_snapshot("user-1", "api-key-1", 150)
        .await
        .expect("auth snapshot should read")
        .expect("auth snapshot should exist");

    let rows = state
        .list_minimal_candidate_selection_rows("openai:chat", "gpt-4.1")
        .await
        .expect("minimal candidate selection rows should read");
    let auth_constraints = SchedulerAuthConstraints {
        allowed_providers: auth_snapshot
            .effective_allowed_providers()
            .map(|items| items.to_vec()),
        allowed_api_formats: auth_snapshot
            .effective_allowed_api_formats()
            .map(|items| items.to_vec()),
        allowed_models: auth_snapshot
            .effective_allowed_models()
            .map(|items| items.to_vec()),
    };

    let selection =
        enumerate_minimal_candidate_selection(EnumerateMinimalCandidateSelectionInput {
            rows,
            normalized_api_format: "openai:chat",
            request_operation: None,
            requested_model_name: "gpt-4.1",
            resolved_global_model_name: "gpt-4.1",
            require_streaming: false,
            required_capabilities: None,
            auth_constraints: Some(&auth_constraints),
        })
        .expect("selection should read");

    assert_eq!(selection.len(), 2);
    assert_eq!(selection[0].provider_id, "provider-1");
    assert_eq!(selection[0].selected_provider_model_name, "gpt-4.1-canary");
    assert_eq!(selection[0].mapping_matched_model, None);
    assert_eq!(selection[1].key_id, "key-3");
    assert_eq!(
        selection[1].selected_provider_model_name,
        "gpt-4.1-edge".to_string()
    );
    assert_eq!(
        selection[1].mapping_matched_model,
        Some("gpt-4.1-edge".to_string())
    );
}

#[tokio::test]
async fn maps_openai_video_task_repository_row_into_read_response() {
    let repository = Arc::new(InMemoryVideoTaskRepository::default());
    repository
        .upsert(UpsertVideoTask {
            id: "task-1".to_string(),
            short_id: Some("short-task-1".to_string()),
            request_id: "request-1".to_string(),
            user_id: Some("user-1".to_string()),
            api_key_id: Some("key-1".to_string()),
            username: Some("user".to_string()),
            api_key_name: Some("primary".to_string()),
            external_task_id: Some("ext-task-1".to_string()),
            provider_id: Some("provider-1".to_string()),
            endpoint_id: Some("endpoint-1".to_string()),
            key_id: Some("provider-key-1".to_string()),
            client_api_format: Some("openai:video".to_string()),
            provider_api_format: Some("openai:video".to_string()),
            format_converted: false,
            model: Some("sora-2".to_string()),
            prompt: Some("hello".to_string()),
            original_request_body: Some(json!({"prompt": "hello"})),
            duration_seconds: Some(4),
            resolution: Some("720p".to_string()),
            aspect_ratio: Some("16:9".to_string()),
            size: Some("1280x720".to_string()),
            status: VideoTaskStatus::Processing,
            progress_percent: 45,
            progress_message: Some("working".to_string()),
            retry_count: 0,
            poll_interval_seconds: 10,
            next_poll_at_unix_secs: Some(120),
            poll_count: 1,
            max_poll_count: 360,
            created_at_unix_ms: 100,
            submitted_at_unix_secs: Some(100),
            completed_at_unix_secs: None,
            updated_at_unix_secs: 120,
            error_code: None,
            error_message: None,
            video_url: None,
            request_metadata: None,
        })
        .await
        .expect("upsert should succeed");

    let state = GatewayDataState::with_video_task_reader_for_tests(repository);
    let response = state
        .read_video_task_response(Some("openai"), "/v1/videos/task-1")
        .await
        .expect("read should succeed")
        .expect("read response should exist");

    assert_eq!(response.status_code, 200);
    assert_eq!(response.body_json["id"], "task-1");
    assert_eq!(response.body_json["status"], "processing");
    assert_eq!(response.body_json["created_at"], 100);
}

#[tokio::test]
async fn maps_gemini_video_task_repository_row_into_read_response() {
    let repository = Arc::new(InMemoryVideoTaskRepository::default());
    repository
        .upsert(UpsertVideoTask {
            id: "task-1".to_string(),
            short_id: Some("localshort123".to_string()),
            request_id: "request-1".to_string(),
            user_id: Some("user-1".to_string()),
            api_key_id: Some("key-1".to_string()),
            username: Some("user".to_string()),
            api_key_name: Some("primary".to_string()),
            external_task_id: Some("operations/ext-task-1".to_string()),
            provider_id: Some("provider-1".to_string()),
            endpoint_id: Some("endpoint-1".to_string()),
            key_id: Some("provider-key-1".to_string()),
            client_api_format: Some("gemini:video".to_string()),
            provider_api_format: Some("gemini:video".to_string()),
            format_converted: false,
            model: Some("veo-3".to_string()),
            prompt: Some("hello".to_string()),
            original_request_body: Some(json!({"prompt": "hello"})),
            duration_seconds: Some(8),
            resolution: Some("720p".to_string()),
            aspect_ratio: Some("16:9".to_string()),
            size: Some("720p".to_string()),
            status: VideoTaskStatus::Completed,
            progress_percent: 100,
            progress_message: None,
            retry_count: 0,
            poll_interval_seconds: 10,
            next_poll_at_unix_secs: None,
            poll_count: 4,
            max_poll_count: 360,
            created_at_unix_ms: 100,
            submitted_at_unix_secs: Some(100),
            completed_at_unix_secs: Some(120),
            updated_at_unix_secs: 120,
            error_code: None,
            error_message: None,
            video_url: None,
            request_metadata: Some(json!({
                "rust_local_snapshot": {
                    "metadata": {
                        "generateVideoResponse": {
                            "generatedSamples": [
                                {
                                    "video": {
                                        "uri": "/v1beta/files/aev_localshort123:download?alt=media"
                                    }
                                }
                            ]
                        }
                    }
                }
            })),
        })
        .await
        .expect("upsert should succeed");

    let state = GatewayDataState::with_video_task_reader_for_tests(repository);
    let response = state
        .read_video_task_response(
            Some("gemini"),
            "/v1beta/models/veo-3/operations/localshort123",
        )
        .await
        .expect("read should succeed")
        .expect("read response should exist");

    assert_eq!(response.status_code, 200);
    assert_eq!(
        response.body_json["name"],
        "models/veo-3/operations/localshort123"
    );
    assert_eq!(response.body_json["done"], true);
}

fn sample_request_candidate(
    id: &str,
    request_id: &str,
    candidate_index: i32,
    status: RequestCandidateStatus,
    started_at_unix_ms: Option<i64>,
    latency_ms: Option<i32>,
    status_code: Option<i32>,
) -> StoredRequestCandidate {
    StoredRequestCandidate::new(
        id.to_string(),
        request_id.to_string(),
        Some("user-1".to_string()),
        Some("api-key-1".to_string()),
        Some("alice".to_string()),
        Some("default".to_string()),
        candidate_index,
        0,
        Some("provider-1".to_string()),
        Some("endpoint-1".to_string()),
        Some("provider-key-1".to_string()),
        status,
        None,
        false,
        status_code,
        None,
        None,
        latency_ms,
        Some(1),
        None,
        None,
        100 + i64::from(candidate_index),
        started_at_unix_ms,
        started_at_unix_ms.map(|value| value + 1),
    )
    .expect("candidate should build")
}

#[tokio::test]
async fn data_state_reads_request_candidate_trace_from_reader() {
    let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
        sample_request_candidate(
            "cand-1",
            "req-1",
            0,
            RequestCandidateStatus::Pending,
            None,
            None,
            None,
        ),
        sample_request_candidate(
            "cand-2",
            "req-1",
            1,
            RequestCandidateStatus::Success,
            Some(101),
            Some(42),
            Some(200),
        ),
    ]));
    let state = GatewayDataState::with_request_candidate_reader_for_tests(repository);

    let trace = state
        .read_request_candidate_trace("req-1", true)
        .await
        .expect("trace should succeed")
        .expect("trace should exist");

    assert_eq!(trace.request_id, "req-1");
    assert_eq!(trace.total_candidates, 1);
    assert_eq!(
        trace.final_status,
        super::candidates::RequestCandidateFinalStatus::Success
    );
    assert_eq!(trace.total_latency_ms, 42);
    assert_eq!(trace.candidates[0].id, "cand-2");
}