aprender-serve 0.64.0

//! API Tests Part 02
//!
//! Generate endpoint tests, streaming tests

#[cfg(feature = "gpu")]
use axum::{
    body::Body,
    http::{Request, StatusCode},
};
#[cfg(feature = "gpu")]
use tower::util::ServiceExt;

#[cfg(feature = "gpu")]
use crate::api::test_helpers::create_test_quantized_model;
#[allow(unused_imports)]
use crate::api::*;

/// IMP-116d: Test cached model can be accessed multiple times (scheduler reuse)
#[tokio::test]
#[cfg(feature = "gpu")]
async fn test_imp_116d_scheduler_reuse_across_requests() {
    use crate::gguf::{GGUFConfig, OwnedQuantizedModelCachedSync};
    use std::sync::Arc;

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = Arc::new(OwnedQuantizedModelCachedSync::new(model));

    // Verify cached model can be accessed multiple times concurrently
    let mut handles = Vec::new();
    for i in 0..5 {
        let model_clone = cached_model.clone();
        handles.push(tokio::spawn(async move {
            // Access model - this exercises the internal scheduler
            let inner = model_clone.model();
            assert_eq!(
                inner.config.hidden_dim, 64,
                "IMP-116d: Access {i} should succeed"
            );
        }));
    }

    // All concurrent accesses should succeed
    for (i, handle) in handles.into_iter().enumerate() {
        handle
            .await
            .unwrap_or_else(|_| panic!("IMP-116d: Concurrent access {i} should not panic"));
    }
}

// ============================================================
// IMP-126: Wire adaptive generation into HTTP serving handler
// RED phase: Tests written first, implementation to follow
// ============================================================

/// IMP-126a: AppState should have dispatch_metrics field
#[test]
#[cfg(feature = "gpu")]
fn test_imp_126a_appstate_has_dispatch_metrics() {
    use crate::gguf::{GGUFConfig, OwnedQuantizedModelCachedSync};

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = OwnedQuantizedModelCachedSync::new(model);

    // Create AppState with cached model
    let state =
        AppState::with_cached_model(cached_model).expect("IMP-126a: Should create AppState");

    // Verify dispatch_metrics is accessible
    let metrics = state.dispatch_metrics();
    assert!(
        metrics.is_some(),
        "IMP-126a: AppState should have dispatch_metrics"
    );

    // Verify metrics starts at zero
    let m = metrics.expect("Should have metrics");
    assert_eq!(
        m.total_dispatches(),
        0,
        "IMP-126a: Metrics should start at zero"
    );
}

/// IMP-126b: OwnedQuantizedModelCachedSync has generate_with_cache_adaptive method
/// This test verifies the method signature exists on the type
#[test]
#[cfg(feature = "gpu")]
fn test_imp_126b_cached_sync_has_generate_adaptive() {
    use crate::gguf::{
        DispatchMetrics, GGUFConfig, OwnedQuantizedModelCachedSync, QuantizedGenerateConfig,
    };
    use std::sync::Arc;

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = OwnedQuantizedModelCachedSync::new(model);
    let metrics = Arc::new(DispatchMetrics::new());

    let gen_config = QuantizedGenerateConfig {
        max_tokens: 3,
        temperature: 0.0,
        top_k: 1,
        stop_tokens: vec![],
        trace: false,
            ..Default::default()
    };

    // Verify method exists by calling it (result may fail due to test model size)
    let prompt = vec![1u32, 2, 3];
    let _result = cached_model.generate_with_cache_adaptive(&prompt, &gen_config, &metrics);

    // IMP-126b: Method exists and can be called
    // Actual generation tested in gguf.rs with proper test model
    assert!(
        true,
        "IMP-126b: generate_with_cache_adaptive method exists on OwnedQuantizedModelCachedSync"
    );
}

/// IMP-126c: AppState provides dispatch_metrics for HTTP handlers
#[test]
#[cfg(feature = "gpu")]
fn test_imp_126c_dispatch_metrics_integration() {
    use crate::gguf::{GGUFConfig, OwnedQuantizedModelCachedSync};
    use std::sync::Arc;

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = OwnedQuantizedModelCachedSync::new(model);

    // Create AppState with cached model - this should initialize dispatch_metrics
    let state =
        AppState::with_cached_model(cached_model).expect("IMP-126c: Should create AppState");

    // Verify dispatch_metrics is accessible and shared
    let metrics1 = state.dispatch_metrics();
    let metrics2 = state.dispatch_metrics();

    assert!(
        metrics1.is_some(),
        "IMP-126c: dispatch_metrics should be available"
    );

    // Metrics should be shareable (Arc)
    let m1 = metrics1.expect("Should have metrics");
    let m2 = metrics2.expect("Should have metrics");
    assert!(
        Arc::ptr_eq(m1, m2),
        "IMP-126c: dispatch_metrics should be shared Arc"
    );
}

/// IMP-126d: Handler uses adaptive generation when dispatch_metrics available
/// This tests that the handler prefers generate_with_cache_adaptive over generate_with_cache
#[test]
#[cfg(feature = "gpu")]
fn test_imp_126d_handler_uses_adaptive_generation() {
    use crate::gguf::{GGUFConfig, OwnedQuantizedModelCachedSync};

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = OwnedQuantizedModelCachedSync::new(model);
    let state =
        AppState::with_cached_model(cached_model).expect("IMP-126d: Should create AppState");

    // Handler should have dispatch_metrics available for adaptive generation
    let metrics = state.dispatch_metrics();
    assert!(
        metrics.is_some(),
        "IMP-126d: Handler should have dispatch_metrics for adaptive generation"
    );

    // Record initial state
    let m = metrics.expect("Should have metrics");
    let initial_cpu = m.cpu_dispatches();
    let initial_gpu = m.gpu_dispatches();

    // The handler code path (test) should use adaptive generation
    // which records dispatches to metrics. We verify the metrics are being
    // passed through by checking they can be incremented.
    m.record_cpu_dispatch();
    m.record_gpu_dispatch();

    assert_eq!(
        m.cpu_dispatches(),
        initial_cpu + 1,
        "IMP-126d: Metrics should track CPU dispatches"
    );
    assert_eq!(
        m.gpu_dispatches(),
        initial_gpu + 1,
        "IMP-126d: Metrics should track GPU dispatches"
    );
}

// ========================================================================
// IMP-127: /metrics/dispatch Endpoint Tests
// ========================================================================

/// IMP-127a: /metrics/dispatch endpoint exists and returns JSON
#[tokio::test]
#[cfg(feature = "gpu")]
async fn test_imp_127a_dispatch_metrics_endpoint_exists() {
    use crate::gguf::{GGUFConfig, OwnedQuantizedModelCachedSync};

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = OwnedQuantizedModelCachedSync::new(model);
    let state =
        AppState::with_cached_model(cached_model).expect("IMP-127a: Should create AppState");

    let app = create_router(state);

    let response = app
        .oneshot(
            Request::builder()
                .uri("/metrics/dispatch")
                .body(Body::empty())
                .expect("test"),
        )
        .await
        .expect("test");

    assert_eq!(
        response.status(),
        StatusCode::OK,
        "IMP-127a: /metrics/dispatch should return 200 OK"
    );

    // Verify JSON content type
    let content_type = response
        .headers()
        .get("content-type")
        .and_then(|v| v.to_str().ok());
    assert!(
        content_type.is_some_and(|s| s.contains("application/json")),
        "IMP-127a: Response should be JSON"
    );
}

/// IMP-127b: /metrics/dispatch returns correct structure
#[tokio::test]
#[cfg(feature = "gpu")]
async fn test_imp_127b_dispatch_metrics_response_structure() {
    use crate::gguf::{GGUFConfig, OwnedQuantizedModelCachedSync};

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = OwnedQuantizedModelCachedSync::new(model);
    let state =
        AppState::with_cached_model(cached_model).expect("IMP-127b: Should create AppState");

    let app = create_router(state);

    let response = app
        .oneshot(
            Request::builder()
                .uri("/metrics/dispatch")
                .body(Body::empty())
                .expect("test"),
        )
        .await
        .expect("test");

    let body = axum::body::to_bytes(response.into_body(), usize::MAX)
        .await
        .expect("test");
    let json: serde_json::Value =
        serde_json::from_slice(&body).expect("IMP-127b: Response should be valid JSON");

    // Verify required fields
    assert!(
        json.get("cpu_dispatches").is_some(),
        "IMP-127b: Response should have cpu_dispatches"
    );
    assert!(
        json.get("gpu_dispatches").is_some(),
        "IMP-127b: Response should have gpu_dispatches"
    );
    assert!(
        json.get("total_dispatches").is_some(),
        "IMP-127b: Response should have total_dispatches"
    );
    assert!(
        json.get("gpu_ratio").is_some(),
        "IMP-127b: Response should have gpu_ratio"
    );
}

/// IMP-127c: /metrics/dispatch starts at zero
#[tokio::test]
#[cfg(feature = "gpu")]
async fn test_imp_127c_dispatch_metrics_starts_zero() {
    use crate::gguf::{GGUFConfig, OwnedQuantizedModelCachedSync};

    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_quantized_model(&config);
    let cached_model = OwnedQuantizedModelCachedSync::new(model);
    let state =
        AppState::with_cached_model(cached_model).expect("IMP-127c: Should create AppState");

    let app = create_router(state);

    let response = app
        .oneshot(
            Request::builder()
                .uri("/metrics/dispatch")
                .body(Body::empty())
                .expect("test"),
        )
        .await
        .expect("test");

    let body = axum::body::to_bytes(response.into_body(), usize::MAX)
        .await
        .expect("test");
    let json: serde_json::Value = serde_json::from_slice(&body).expect("test");

    assert_eq!(
        json["cpu_dispatches"].as_u64(),
        Some(0),
        "IMP-127c: cpu_dispatches should start at 0"
    );
    assert_eq!(
        json["gpu_dispatches"].as_u64(),
        Some(0),
        "IMP-127c: gpu_dispatches should start at 0"
    );
    assert_eq!(
        json["total_dispatches"].as_u64(),
        Some(0),
        "IMP-127c: total_dispatches should start at 0"
    );
}

include!("imp_127d.rs");
include!("imp_130d.rs");
include!("imp_132a.rs");