foxhunt/services/trading_service/tests/ensemble_audit_tests.rs

#![allow(unexpected_cfgs)]
#![cfg(feature = "__trading_service_integration")]
//! Integration tests for ensemble audit logging
//!
//! These tests validate the PostgreSQL audit logging system for ensemble predictions,
//! including prediction inserts, P&L updates, and analysis queries.

use sqlx::PgPool;
use std::collections::HashMap;
use uuid::Uuid;

// Import ensemble types from ml crate
use ml::ensemble::{EnsembleDecision, ModelVote, TradingAction};

// Test database URL
fn get_test_db_url() -> String {
    std::env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
    })
}

#[tokio::test]
async fn test_ensemble_audit_logger_initialization() {
    // Create database pool
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Import audit logger (we'll need to add the module reference)
    // For now, we'll test the SQL schema directly

    // Verify tables exist
    let result = sqlx::query!(
        r#"
        SELECT tablename
        FROM pg_tables
        WHERE schemaname = 'public'
          AND tablename IN ('ensemble_predictions', 'model_performance_attribution', 'ab_test_experiments')
        "#
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to query tables");

    assert_eq!(result.len(), 3, "All ensemble audit tables should exist");

    pool.close().await;
}

#[tokio::test]
async fn test_log_ensemble_prediction() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Create mock ensemble decision
    let mut model_votes = HashMap::new();
    model_votes.insert(
        "DQN".to_string(),
        ModelVote::new("DQN".to_string(), 0.8, 0.9, 0.33),
    );
    model_votes.insert(
        "PPO".to_string(),
        ModelVote::new("PPO".to_string(), 0.7, 0.85, 0.33),
    );
    model_votes.insert(
        "TFT".to_string(),
        ModelVote::new("TFT".to_string(), 0.6, 0.8, 0.34),
    );

    let decision = EnsembleDecision::new(TradingAction::Buy, 0.85, 0.7, 0.15, model_votes);

    // Insert prediction manually (simulating audit logger)
    let id = Uuid::new_v4();
    let symbol = "ES.FUT";

    let result = sqlx::query!(
        r#"
        INSERT INTO ensemble_predictions (
            id, symbol, ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate,
            dqn_signal, dqn_confidence, dqn_weight, dqn_vote,
            ppo_signal, ppo_confidence, ppo_weight, ppo_vote,
            tft_signal, tft_confidence, tft_weight, tft_vote
        ) VALUES (
            $1, $2, $3, $4, $5, $6,
            $7, $8, $9, $10,
            $11, $12, $13, $14,
            $15, $16, $17, $18
        )
        "#,
        id,
        symbol,
        "BUY",
        0.7,
        0.85,
        0.15,
        0.8,
        0.9,
        0.33,
        "BUY",
        0.7,
        0.85,
        0.33,
        "BUY",
        0.6,
        0.8,
        0.34,
        "BUY",
    )
    .execute(&pool)
    .await;

    assert!(result.is_ok(), "Prediction insert should succeed");

    // Verify prediction was inserted
    let record = sqlx::query!(
        r#"
        SELECT id, symbol, ensemble_action, ensemble_confidence, disagreement_rate
        FROM ensemble_predictions
        WHERE id = $1
        "#,
        id
    )
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch prediction");

    assert_eq!(record.symbol, "ES.FUT");
    assert_eq!(record.ensemble_action, "BUY");
    assert!((record.ensemble_confidence - 0.85).abs() < 0.001);
    assert!((record.disagreement_rate - 0.15).abs() < 0.001);

    // Cleanup
    sqlx::query!("DELETE FROM ensemble_predictions WHERE id = $1", id)
        .execute(&pool)
        .await
        .expect("Failed to cleanup");

    pool.close().await;
}

#[tokio::test]
async fn test_update_prediction_pnl() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Insert a test prediction
    let id = Uuid::new_v4();

    sqlx::query!(
        r#"
        INSERT INTO ensemble_predictions (
            id, symbol, ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate
        ) VALUES (
            $1, 'NQ.FUT', 'BUY', 0.5, 0.7, 0.2
        )
        "#,
        id
    )
    .execute(&pool)
    .await
    .expect("Failed to insert test prediction");

    // Update P&L (simulating trade close)
    let pnl_cents = 1250; // $12.50 profit
    let commission_cents = 25; // $0.25 commission
    let slippage_bps = 5; // 5 basis points

    let result = sqlx::query!(
        r#"
        UPDATE ensemble_predictions
        SET pnl = $2, commission = $3, slippage_bps = $4
        WHERE id = $1
        "#,
        id,
        pnl_cents,
        commission_cents,
        slippage_bps,
    )
    .execute(&pool)
    .await;

    assert!(result.is_ok(), "P&L update should succeed");

    // Verify update
    let record = sqlx::query!(
        r#"
        SELECT pnl, commission, slippage_bps
        FROM ensemble_predictions
        WHERE id = $1
        "#,
        id
    )
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch updated prediction");

    assert_eq!(record.pnl, Some(pnl_cents));
    assert_eq!(record.commission, Some(commission_cents));
    assert_eq!(record.slippage_bps, Some(slippage_bps));

    // Cleanup
    sqlx::query!("DELETE FROM ensemble_predictions WHERE id = $1", id)
        .execute(&pool)
        .await
        .expect("Failed to cleanup");

    pool.close().await;
}

#[tokio::test]
async fn test_model_performance_attribution() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Insert model performance record
    let id = Uuid::new_v4();
    let model_id = "DQN";
    let symbol = "ES.FUT";
    let window_hours = 24;
    let total_predictions = 100;
    let correct_predictions = 58;
    let accuracy = 0.58;
    let total_pnl_cents = 5420; // $54.20
    let sharpe_ratio = 1.85;
    let avg_weight = 0.33;
    let avg_confidence = 0.82;

    let result = sqlx::query!(
        r#"
        INSERT INTO model_performance_attribution (
            id, model_id, symbol, window_hours,
            total_predictions, correct_predictions, accuracy,
            total_pnl, sharpe_ratio,
            avg_weight, avg_confidence
        ) VALUES (
            $1, $2, $3, $4,
            $5, $6, $7,
            $8, $9,
            $10, $11
        )
        "#,
        id,
        model_id,
        symbol,
        window_hours,
        total_predictions,
        correct_predictions,
        accuracy,
        total_pnl_cents as i64,
        sharpe_ratio,
        avg_weight,
        avg_confidence,
    )
    .execute(&pool)
    .await;

    assert!(
        result.is_ok(),
        "Performance attribution insert should succeed"
    );

    // Verify insertion
    let record = sqlx::query!(
        r#"
        SELECT model_id, symbol, accuracy, sharpe_ratio, total_pnl
        FROM model_performance_attribution
        WHERE id = $1
        "#,
        id
    )
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch performance record");

    assert_eq!(record.model_id, "DQN");
    assert_eq!(record.symbol, "ES.FUT");
    assert!((record.accuracy - 0.58).abs() < 0.001);
    assert_eq!(record.total_pnl, total_pnl_cents as i64);

    // Cleanup
    sqlx::query!(
        "DELETE FROM model_performance_attribution WHERE id = $1",
        id
    )
    .execute(&pool)
    .await
    .expect("Failed to cleanup");

    pool.close().await;
}

#[tokio::test]
async fn test_ab_test_experiment_tracking() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Create A/B test experiment
    let test_id = Uuid::new_v4();
    let test_name = format!("test_ensemble_vs_dqn_{}", test_id);

    let result = sqlx::query!(
        r#"
        INSERT INTO ab_test_experiments (
            test_id, test_name, description,
            control_variant, treatment_variant,
            traffic_split, min_sample_size, significance_level,
            status
        ) VALUES (
            $1, $2, 'Testing ensemble vs single DQN model',
            'DQN_ONLY', 'ENSEMBLE',
            0.5, 1000, 0.05,
            'running'
        )
        "#,
        test_id,
        test_name,
    )
    .execute(&pool)
    .await;

    assert!(result.is_ok(), "A/B test creation should succeed");

    // Insert predictions for control group
    for i in 0..5 {
        let pred_id = Uuid::new_v4();
        sqlx::query!(
            r#"
            INSERT INTO ensemble_predictions (
                id, symbol, ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate,
                ab_test_id, ab_group
            ) VALUES (
                $1, 'ES.FUT', 'BUY', 0.6, 0.75, 0.2,
                $2, 'control'
            )
            "#,
            pred_id,
            test_id,
        )
        .execute(&pool)
        .await
        .expect("Failed to insert control prediction");
    }

    // Insert predictions for treatment group
    for i in 0..5 {
        let pred_id = Uuid::new_v4();
        sqlx::query!(
            r#"
            INSERT INTO ensemble_predictions (
                id, symbol, ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate,
                ab_test_id, ab_group
            ) VALUES (
                $1, 'ES.FUT', 'BUY', 0.7, 0.85, 0.15,
                $2, 'treatment'
            )
            "#,
            pred_id,
            test_id,
        )
        .execute(&pool)
        .await
        .expect("Failed to insert treatment prediction");
    }

    // Query A/B test assignments
    let results = sqlx::query!(
        r#"
        SELECT ab_group, COUNT(*) as count
        FROM ensemble_predictions
        WHERE ab_test_id = $1
        GROUP BY ab_group
        "#,
        test_id
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to query A/B test results");

    assert_eq!(
        results.len(),
        2,
        "Should have both control and treatment groups"
    );

    for record in results {
        let count = record.count.expect("Count should not be null");
        assert_eq!(count, 5, "Each group should have 5 predictions");
    }

    // Cleanup
    sqlx::query!(
        "DELETE FROM ensemble_predictions WHERE ab_test_id = $1",
        test_id
    )
    .execute(&pool)
    .await
    .expect("Failed to cleanup predictions");

    sqlx::query!(
        "DELETE FROM ab_test_experiments WHERE test_id = $1",
        test_id
    )
    .execute(&pool)
    .await
    .expect("Failed to cleanup experiment");

    pool.close().await;
}

#[tokio::test]
async fn test_batch_prediction_insert_performance() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    let batch_size = 100;
    let start = std::time::Instant::now();

    // Insert 100 predictions in a transaction
    let mut tx = pool.begin().await.expect("Failed to start transaction");

    for i in 0..batch_size {
        let id = Uuid::new_v4();
        sqlx::query!(
            r#"
            INSERT INTO ensemble_predictions (
                id, symbol, ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate,
                inference_latency_us, aggregation_latency_us
            ) VALUES (
                $1, 'ES.FUT', 'BUY', 0.6, 0.75, 0.2,
                42, 8
            )
            "#,
            id
        )
        .execute(&mut *tx)
        .await
        .expect("Failed to insert prediction in batch");
    }

    tx.commit().await.expect("Failed to commit transaction");

    let elapsed = start.elapsed();
    let throughput = batch_size as f64 / elapsed.as_secs_f64();

    println!(
        "Batch insert performance: {} predictions in {:.2}ms ({:.0} predictions/sec)",
        batch_size,
        elapsed.as_secs_f64() * 1000.0,
        throughput
    );

    // Performance assertion: Should handle >1000 predictions/sec
    assert!(
        throughput > 100.0,
        "Batch insert throughput should exceed 100 predictions/sec, got {:.0}",
        throughput
    );

    // Cleanup
    sqlx::query!(
        "DELETE FROM ensemble_predictions WHERE symbol = 'ES.FUT' AND inference_latency_us = 42"
    )
    .execute(&pool)
    .await
    .expect("Failed to cleanup batch");

    pool.close().await;
}

#[tokio::test]
async fn test_analysis_query_performance() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Test high disagreement query performance
    let start = std::time::Instant::now();

    let _results = sqlx::query!(
        r#"
        SELECT id, symbol, disagreement_rate, ensemble_action
        FROM ensemble_predictions
        WHERE disagreement_rate > 0.5
        ORDER BY disagreement_rate DESC
        LIMIT 100
        "#
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to execute high disagreement query");

    let elapsed = start.elapsed();

    println!(
        "High disagreement query latency: {:.2}ms",
        elapsed.as_secs_f64() * 1000.0
    );

    // Performance assertion: Query should complete in <100ms
    assert!(
        elapsed.as_millis() < 100,
        "High disagreement query should complete in <100ms, took {}ms",
        elapsed.as_millis()
    );

    pool.close().await;
}

#[tokio::test]
async fn test_timescaledb_hypertable_functionality() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Verify hypertable was created for ensemble_predictions
    let result = sqlx::query!(
        r#"
        SELECT hypertable_name
        FROM timescaledb_information.hypertable
        WHERE hypertable_name IN ('ensemble_predictions', 'model_performance_attribution')
        "#
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to query hypertables");

    assert_eq!(
        result.len(),
        2,
        "Both ensemble_predictions and model_performance_attribution should be hypertables"
    );

    pool.close().await;
}

#[tokio::test]
async fn test_feature_snapshot_jsonb() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Create feature snapshot JSON
    let feature_snapshot = serde_json::json!({
        "ohlcv": {
            "open": 5000.50,
            "high": 5010.25,
            "low": 4995.00,
            "close": 5005.75,
            "volume": 125000
        },
        "technical_indicators": {
            "rsi": 62.5,
            "macd": 12.3,
            "bollinger_upper": 5020.0,
            "bollinger_lower": 4990.0
        }
    });

    let id = Uuid::new_v4();

    // Insert prediction with feature snapshot
    let result = sqlx::query!(
        r#"
        INSERT INTO ensemble_predictions (
            id, symbol, ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate,
            feature_snapshot
        ) VALUES (
            $1, 'ES.FUT', 'BUY', 0.6, 0.75, 0.2,
            $2
        )
        "#,
        id,
        feature_snapshot,
    )
    .execute(&pool)
    .await;

    assert!(result.is_ok(), "Feature snapshot insert should succeed");

    // Query feature snapshot
    let record = sqlx::query!(
        r#"
        SELECT feature_snapshot
        FROM ensemble_predictions
        WHERE id = $1
        "#,
        id
    )
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch feature snapshot");

    let snapshot = record
        .feature_snapshot
        .expect("Feature snapshot should not be null");
    let rsi = snapshot["technical_indicators"]["rsi"].as_f64();
    assert_eq!(rsi, Some(62.5), "RSI should match");

    // Cleanup
    sqlx::query!("DELETE FROM ensemble_predictions WHERE id = $1", id)
        .execute(&pool)
        .await
        .expect("Failed to cleanup");

    pool.close().await;
}

#[tokio::test]
async fn test_continuous_aggregate_views() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Query hourly ensemble performance view
    let results = sqlx::query!(
        r#"
        SELECT bucket, symbol, prediction_count, avg_confidence
        FROM ensemble_performance_hourly
        ORDER BY bucket DESC
        LIMIT 10
        "#
    )
    .fetch_all(&pool)
    .await;

    // View should exist (even if empty)
    assert!(
        results.is_ok(),
        "Continuous aggregate view should be queryable"
    );

    // Query daily model performance view
    let results = sqlx::query!(
        r#"
        SELECT bucket, model_id, avg_accuracy, avg_sharpe_ratio
        FROM model_performance_daily
        ORDER BY bucket DESC
        LIMIT 10
        "#
    )
    .fetch_all(&pool)
    .await;

    assert!(
        results.is_ok(),
        "Daily model performance view should be queryable"
    );

    pool.close().await;
}

#[tokio::test]
async fn test_utility_functions() {
    let pool = PgPool::connect(&get_test_db_url())
        .await
        .expect("Failed to connect to test database");

    // Test get_top_models_24h function
    let results = sqlx::query!(
        r#"
        SELECT * FROM get_top_models_24h(NULL::VARCHAR, 5)
        "#
    )
    .fetch_all(&pool)
    .await;

    assert!(
        results.is_ok(),
        "get_top_models_24h function should execute"
    );

    // Test calculate_model_correlation_7d function
    let results = sqlx::query!(
        r#"
        SELECT * FROM calculate_model_correlation_7d(NULL::VARCHAR)
        "#
    )
    .fetch_all(&pool)
    .await;

    assert!(
        results.is_ok(),
        "calculate_model_correlation_7d function should execute"
    );

    // Test get_high_disagreement_events_24h function
    let results = sqlx::query!(
        r#"
        SELECT * FROM get_high_disagreement_events_24h(NULL::VARCHAR, 0.5, 100)
        "#
    )
    .fetch_all(&pool)
    .await;

    assert!(
        results.is_ok(),
        "get_high_disagreement_events_24h function should execute"
    );

    pool.close().await;
}