foxhunt/services/ml_training_service/tests/integration/database_integration_test.rs

//! Database Integration Tests
//!
//! Tests migration 046 schema validation, foreign key cascades, triggers, and indexes
//!
//! **Test Coverage:**
//! - Migration 046 schema validation
//! - Parent-child foreign key cascade
//! - Trigger-based batch aggregation
//! - Concurrent transaction isolation
//! - Index performance validation (<10ms queries)

use std::time::Instant;

use sqlx::PgPool;
use std::time::Duration;
use uuid::Uuid;

/// Create test database pool
async fn create_test_pool() -> PgPool {
    let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
    });

    PgPool::connect(&database_url)
        .await
        .expect("Failed to connect to database")
}

/// Helper to create a test batch job
async fn create_test_batch(pool: &PgPool, name: &str) -> Uuid {
    let id = Uuid::new_v4();

    sqlx::query(
        "INSERT INTO batch_jobs (id, name, description, status, config_json)
         VALUES ($1, $2, $3, 'Pending', '{}'::jsonb)"
    )
    .bind(id)
    .bind(name)
    .bind("Test batch job")
    .execute(pool)
    .await
    .expect("Failed to create batch job");

    id
}

/// Helper to create a test child job
async fn create_test_child(pool: &PgPool, batch_id: Uuid, model_type: &str) -> Uuid {
    let id = Uuid::new_v4();

    sqlx::query(
        "INSERT INTO child_jobs (id, batch_id, model_type, status, total_epochs, config_json)
         VALUES ($1, $2, $3, 'Pending', 10, '{}'::jsonb)"
    )
    .bind(id)
    .bind(batch_id)
    .bind(model_type)
    .execute(pool)
    .await
    .expect("Failed to create child job");

    id
}

/// Clean up test data
async fn cleanup_test_data(pool: &PgPool, batch_id: Uuid) {
    // Delete batch (cascade should handle children)
    sqlx::query("DELETE FROM batch_jobs WHERE id = $1")
        .bind(batch_id)
        .execute(pool)
        .await
        .expect("Failed to cleanup batch");
}

// ============================================================================
// Test 1: Migration 046 schema validation
// ============================================================================

#[tokio::test]
async fn test_migration_046_schema() {
    let pool = create_test_pool().await;

    // Verify batch_jobs table exists
    let batch_table_exists: bool = sqlx::query_scalar(
        "SELECT EXISTS (
            SELECT FROM information_schema.tables
            WHERE table_name = 'batch_jobs'
        )"
    )
    .fetch_one(&pool)
    .await
    .expect("Failed to check batch_jobs table");

    assert!(batch_table_exists, "batch_jobs table does not exist");
    println!("✅ batch_jobs table exists");

    // Verify child_jobs table exists
    let child_table_exists: bool = sqlx::query_scalar(
        "SELECT EXISTS (
            SELECT FROM information_schema.tables
            WHERE table_name = 'child_jobs'
        )"
    )
    .fetch_one(&pool)
    .await
    .expect("Failed to check child_jobs table");

    assert!(child_table_exists, "child_jobs table does not exist");
    println!("✅ child_jobs table exists");

    // Verify batch_jobs columns
    let batch_columns: Vec<String> = sqlx::query_scalar(
        "SELECT column_name FROM information_schema.columns
         WHERE table_name = 'batch_jobs'
         ORDER BY column_name"
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to fetch batch_jobs columns");

    let expected_batch_columns = vec![
        "completed_at", "completed_jobs", "config_json", "created_at",
        "description", "failed_jobs", "id", "name", "overall_progress",
        "pending_jobs", "running_jobs", "started_at", "status", "total_jobs"
    ];

    for col in &expected_batch_columns {
        assert!(
            batch_columns.contains(&col.to_string()),
            "Missing batch_jobs column: {}",
            col
        );
    }

    println!("✅ All batch_jobs columns present: {}", batch_columns.len());

    // Verify child_jobs columns
    let child_columns: Vec<String> = sqlx::query_scalar(
        "SELECT column_name FROM information_schema.columns
         WHERE table_name = 'child_jobs'
         ORDER BY column_name"
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to fetch child_jobs columns");

    let expected_child_columns = vec![
        "batch_id", "completed_at", "config_json", "created_at", "current_epoch",
        "error_message", "id", "model_type", "model_weight", "progress_pct",
        "started_at", "status", "total_epochs", "updated_at"
    ];

    for col in &expected_child_columns {
        assert!(
            child_columns.contains(&col.to_string()),
            "Missing child_jobs column: {}",
            col
        );
    }

    println!("✅ All child_jobs columns present: {}", child_columns.len());

    // Verify constraints
    let constraints: Vec<String> = sqlx::query_scalar(
        "SELECT constraint_name FROM information_schema.table_constraints
         WHERE table_name IN ('batch_jobs', 'child_jobs')
         AND constraint_type = 'CHECK'"
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to fetch constraints");

    assert!(constraints.len() >= 5, "Expected at least 5 CHECK constraints");
    println!("✅ CHECK constraints validated: {}", constraints.len());

    pool.close().await;
}

// ============================================================================
// Test 2: Parent-child foreign key cascade
// ============================================================================

#[tokio::test]
async fn test_foreign_key_cascade() {
    let pool = create_test_pool().await;

    // Create a batch with children
    let batch_id = create_test_batch(&pool, "cascade_test").await;
    let child1_id = create_test_child(&pool, batch_id, "DQN").await;
    let child2_id = create_test_child(&pool, batch_id, "PPO").await;

    println!("✅ Created batch {} with 2 children", batch_id);

    // Verify children exist
    let child_count: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM child_jobs WHERE batch_id = $1"
    )
    .bind(batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to count children");

    assert_eq!(child_count, 2);
    println!("✅ Verified 2 children exist");

    // Delete parent (should cascade to children)
    sqlx::query("DELETE FROM batch_jobs WHERE id = $1")
        .bind(batch_id)
        .execute(&pool)
        .await
        .expect("Failed to delete batch");

    println!("✅ Deleted parent batch");

    // Verify children are gone (cascade)
    let remaining_children: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM child_jobs WHERE batch_id = $1"
    )
    .bind(batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to count remaining children");

    assert_eq!(remaining_children, 0);
    println!("✅ CASCADE delete validated: children removed");

    // Verify specific children don't exist
    let child1_exists: bool = sqlx::query_scalar(
        "SELECT EXISTS (SELECT 1 FROM child_jobs WHERE id = $1)"
    )
    .bind(child1_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to check child1");

    let child2_exists: bool = sqlx::query_scalar(
        "SELECT EXISTS (SELECT 1 FROM child_jobs WHERE id = $1)"
    )
    .bind(child2_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to check child2");

    assert!(!child1_exists && !child2_exists);
    println!("✅ Both children cascade-deleted successfully");

    pool.close().await;
}

// ============================================================================
// Test 3: Trigger-based batch aggregation
// ============================================================================

#[tokio::test]
async fn test_trigger_batch_aggregation() {
    let pool = create_test_pool().await;

    // Create batch
    let batch_id = create_test_batch(&pool, "aggregation_test").await;

    // Create 4 children (DQN, PPO, MAMBA-2, TFT)
    let models = vec!["DQN", "PPO", "MAMBA-2", "TFT"];
    for model in &models {
        create_test_child(&pool, batch_id, model).await;
    }

    println!("✅ Created batch with 4 children");

    // Verify batch aggregates (trigger should update)
    let batch: (i32, i32, i32, i32, i32, f64) = sqlx::query_as(
        "SELECT total_jobs, pending_jobs, running_jobs, completed_jobs, failed_jobs, overall_progress
         FROM batch_jobs WHERE id = $1"
    )
    .bind(batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch batch");

    assert_eq!(batch.0, 4); // total_jobs
    assert_eq!(batch.1, 4); // pending_jobs
    assert_eq!(batch.2, 0); // running_jobs
    assert_eq!(batch.3, 0); // completed_jobs
    assert_eq!(batch.4, 0); // failed_jobs
    println!("✅ Initial aggregates: total=4, pending=4");

    // Update one child to Running
    sqlx::query(
        "UPDATE child_jobs
         SET status = 'Running', progress_pct = 50.0
         WHERE batch_id = $1 AND model_type = 'DQN'"
    )
    .bind(batch_id)
    .execute(&pool)
    .await
    .expect("Failed to update child");

    // Verify trigger updated batch
    let batch_after: (String, i32, i32, f64) = sqlx::query_as(
        "SELECT status, pending_jobs, running_jobs, overall_progress
         FROM batch_jobs WHERE id = $1"
    )
    .bind(batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch batch");

    assert_eq!(batch_after.0, "Running"); // status
    assert_eq!(batch_after.1, 3); // pending_jobs (3 left)
    assert_eq!(batch_after.2, 1); // running_jobs (1 running)
    assert!(batch_after.3 > 0.0); // overall_progress > 0
    println!("✅ Trigger updated batch: status=Running, progress={:.1}%", batch_after.3);

    // Complete all children
    sqlx::query(
        "UPDATE child_jobs
         SET status = 'Completed', progress_pct = 100.0
         WHERE batch_id = $1"
    )
    .bind(batch_id)
    .execute(&pool)
    .await
    .expect("Failed to complete children");

    // Verify batch is Completed
    let final_batch: (String, i32, f64) = sqlx::query_as(
        "SELECT status, completed_jobs, overall_progress
         FROM batch_jobs WHERE id = $1"
    )
    .bind(batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to fetch final batch");

    assert_eq!(final_batch.0, "Completed");
    assert_eq!(final_batch.1, 4); // all completed
    assert_eq!(final_batch.2, 100.0); // 100% progress
    println!("✅ Batch completion detected: status=Completed, progress=100.0%");

    cleanup_test_data(&pool, batch_id).await;
    pool.close().await;
}

// ============================================================================
// Test 4: Concurrent transaction isolation
// ============================================================================

#[tokio::test]
async fn test_concurrent_transactions() {
    let pool = create_test_pool().await;

    // Create batch
    let batch_id = create_test_batch(&pool, "concurrent_test").await;

    // Create 2 children
    let child1_id = create_test_child(&pool, batch_id, "DQN").await;
    let child2_id = create_test_child(&pool, batch_id, "PPO").await;

    println!("✅ Created batch with 2 children");

    // Spawn concurrent updates
    let pool1 = pool.clone();
    let pool2 = pool.clone();

    let handle1 = tokio::spawn(async move {
        // Transaction 1: Update child1
        let mut tx = pool1.begin().await.expect("Failed to begin tx1");

        sqlx::query(
            "UPDATE child_jobs
             SET status = 'Running', progress_pct = 25.0
             WHERE id = $1"
        )
        .bind(child1_id)
        .execute(&mut *tx)
        .await
        .expect("Failed to update child1");

        // Small delay to test isolation
        tokio::time::sleep(Duration::from_millis(10)).await;

        tx.commit().await.expect("Failed to commit tx1");
    });

    let handle2 = tokio::spawn(async move {
        // Transaction 2: Update child2
        let mut tx = pool2.begin().await.expect("Failed to begin tx2");

        sqlx::query(
            "UPDATE child_jobs
             SET status = 'Running', progress_pct = 75.0
             WHERE id = $1"
        )
        .bind(child2_id)
        .execute(&mut *tx)
        .await
        .expect("Failed to update child2");

        // Small delay to test isolation
        tokio::time::sleep(Duration::from_millis(10)).await;

        tx.commit().await.expect("Failed to commit tx2");
    });

    // Wait for both transactions
    handle1.await.expect("Transaction 1 panicked");
    handle2.await.expect("Transaction 2 panicked");

    println!("✅ Concurrent transactions completed");

    // Verify both updates succeeded
    let children: Vec<(String, f64)> = sqlx::query_as(
        "SELECT status, progress_pct FROM child_jobs
         WHERE batch_id = $1
         ORDER BY model_type"
    )
    .bind(batch_id)
    .fetch_all(&pool)
    .await
    .expect("Failed to fetch children");

    assert_eq!(children.len(), 2);
    assert_eq!(children[0].0, "Running"); // DQN
    assert_eq!(children[0].1, 25.0);
    assert_eq!(children[1].0, "Running"); // PPO
    assert_eq!(children[1].1, 75.0);

    println!("✅ Transaction isolation validated: both updates persisted");

    cleanup_test_data(&pool, batch_id).await;
    pool.close().await;
}

// ============================================================================
// Test 5: Index performance validation (<10ms queries)
// ============================================================================

#[tokio::test]
async fn test_index_performance() {
    let pool = create_test_pool().await;

    // Create test data
    let batch_id = create_test_batch(&pool, "perf_test").await;
    for model in &["DQN", "PPO", "MAMBA-2", "TFT"] {
        create_test_child(&pool, batch_id, model).await;
    }

    println!("✅ Created test data");

    // Test 1: Query by batch status (should use idx_batch_jobs_status)
    let start = Instant::now();
    let _batches: Vec<String> = sqlx::query_scalar(
        "SELECT name FROM batch_jobs WHERE status = 'Pending'"
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to query batches");
    let duration1 = start.elapsed();

    assert!(duration1 < Duration::from_millis(10), "Query 1 too slow: {:?}", duration1);
    println!("✅ Batch status query: {:?} (<10ms)", duration1);

    // Test 2: Query by child batch_id (should use idx_child_jobs_batch_id)
    let start = Instant::now();
    let _children: Vec<String> = sqlx::query_scalar(
        "SELECT model_type FROM child_jobs WHERE batch_id = $1"
    )
    .bind(batch_id)
    .fetch_all(&pool)
    .await
    .expect("Failed to query children");
    let duration2 = start.elapsed();

    assert!(duration2 < Duration::from_millis(10), "Query 2 too slow: {:?}", duration2);
    println!("✅ Child batch_id query: {:?} (<10ms)", duration2);

    // Test 3: Composite query (should use idx_child_jobs_batch_status)
    let start = Instant::now();
    let _count: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM child_jobs
         WHERE batch_id = $1 AND status = 'Pending'"
    )
    .bind(batch_id)
    .fetch_one(&pool)
    .await
    .expect("Failed to count children");
    let duration3 = start.elapsed();

    assert!(duration3 < Duration::from_millis(10), "Query 3 too slow: {:?}", duration3);
    println!("✅ Composite batch+status query: {:?} (<10ms)", duration3);

    // Test 4: JSONB query (should use GIN index)
    let start = Instant::now();
    let _batches: Vec<String> = sqlx::query_scalar(
        "SELECT name FROM batch_jobs
         WHERE config_json @> '{}'::jsonb"
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to query JSONB");
    let duration4 = start.elapsed();

    assert!(duration4 < Duration::from_millis(10), "Query 4 too slow: {:?}", duration4);
    println!("✅ JSONB GIN index query: {:?} (<10ms)", duration4);

    // Test 5: ORDER BY created_at (should use idx_batch_jobs_created_at)
    let start = Instant::now();
    let _batches: Vec<String> = sqlx::query_scalar(
        "SELECT name FROM batch_jobs
         ORDER BY created_at DESC
         LIMIT 10"
    )
    .fetch_all(&pool)
    .await
    .expect("Failed to query ordered batches");
    let duration5 = start.elapsed();

    assert!(duration5 < Duration::from_millis(10), "Query 5 too slow: {:?}", duration5);
    println!("✅ ORDER BY created_at query: {:?} (<10ms)", duration5);

    println!("✅ All index performance tests passed (<10ms)");

    cleanup_test_data(&pool, batch_id).await;
    pool.close().await;
}