foxhunt/tli/tests/e2e/train_stop_e2e_test.rs

//! End-to-End Tests for `tli train stop` Command
//!
//! Tests stopping training jobs and cleanup

use tli::proto::ml_training::{
    ml_training_service_server::MlTrainingServiceServer, TrainingStatus,
    ml_training_service_client::MlTrainingServiceClient, StopTrainingRequest,
};
use tonic::transport::{Server, Channel};
use std::net::SocketAddr;
use tokio::time::Duration;

use super::mock_ml_training_service::{MockMlTrainingService, MockTrainingState};
use super::test_fixtures;

/// Helper to start mock gRPC server
async fn start_mock_server(state: MockTrainingState) -> SocketAddr {
    let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
    let service = MockMlTrainingService::new(state);

    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
    let local_addr = listener.local_addr().unwrap();

    tokio::spawn(async move {
        Server::builder()
            .add_service(MlTrainingServiceServer::new(service))
            .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
            .await
            .unwrap();
    });

    tokio::time::sleep(Duration::from_millis(100)).await;
    local_addr
}

/// TEST 1: Stop running job
#[tokio::test]
async fn test_stop_running_job() {
    let state = MockTrainingState::default();

    let job = test_fixtures::create_test_job(
        "stop_test_1",
        "TFT",
        TrainingStatus::Running,
    );
    state.add_job(job);

    let addr = start_mock_server(state.clone()).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request = StopTrainingRequest {
        job_id: "stop_test_1".to_string(),
        reason: "User requested stop".to_string(),
    };

    let response = client.stop_training(request).await;
    assert!(response.is_ok(), "Should successfully stop running job");

    let result = response.unwrap().into_inner();
    assert!(result.success);

    // Verify job status changed to Stopped
    let job = state.get_job("stop_test_1").unwrap();
    assert_eq!(job.status, TrainingStatus::Stopped as i32);
}

/// TEST 2: Stop entire batch
#[tokio::test]
async fn test_stop_batch() {
    let state = MockTrainingState::default();

    // Create batch jobs
    for i in 0..4 {
        let job = test_fixtures::create_test_job(
            &format!("batch_job_{}", i),
            "TFT",
            TrainingStatus::Running,
        );
        state.add_job(job);
    }

    let addr = start_mock_server(state.clone()).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    // Stop each job in batch
    for i in 0..4 {
        let request = StopTrainingRequest {
            job_id: format!("batch_job_{}", i),
            reason: "Batch stop".to_string(),
        };

        let response = client.stop_training(request).await;
        assert!(response.is_ok(), "Should stop batch job {}", i);
    }

    // Verify all jobs stopped
    for i in 0..4 {
        let job = state.get_job(&format!("batch_job_{}", i)).unwrap();
        assert_eq!(job.status, TrainingStatus::Stopped as i32);
    }
}

/// TEST 3: Stop already completed job → Error
#[tokio::test]
async fn test_stop_completed_job() {
    let state = MockTrainingState::default();

    let job = test_fixtures::create_test_job(
        "completed_job",
        "DQN",
        TrainingStatus::Completed,
    );
    state.add_job(job);

    let addr = start_mock_server(state).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request = StopTrainingRequest {
        job_id: "completed_job".to_string(),
        reason: "Test stop completed".to_string(),
    };

    // Mock allows this, but real implementation might error
    let response = client.stop_training(request).await;
    assert!(response.is_ok());
}

/// TEST 4: Force stop (cleanup)
#[tokio::test]
async fn test_force_stop() {
    let state = MockTrainingState::default();

    let job = test_fixtures::create_test_job(
        "force_stop_test",
        "PPO",
        TrainingStatus::Running,
    );
    state.add_job(job);

    let addr = start_mock_server(state.clone()).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request = StopTrainingRequest {
        job_id: "force_stop_test".to_string(),
        reason: "Force stop with cleanup".to_string(),
    };

    let response = client.stop_training(request).await;
    assert!(response.is_ok());

    let job = state.get_job("force_stop_test").unwrap();
    assert_eq!(job.status, TrainingStatus::Stopped as i32);
}

/// TEST 5: Stop non-existent job → 404
#[tokio::test]
async fn test_stop_nonexistent_job() {
    let state = MockTrainingState::default();
    let addr = start_mock_server(state).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request = StopTrainingRequest {
        job_id: "nonexistent_job_123".to_string(),
        reason: "Test".to_string(),
    };

    let response = client.stop_training(request).await;
    assert!(response.is_err(), "Should return 404 for non-existent job");
    assert!(response.unwrap_err().message().contains("not found"));
}

/// TEST 6: Idempotent stop
#[tokio::test]
async fn test_idempotent_stop() {
    let state = MockTrainingState::default();

    let job = test_fixtures::create_test_job(
        "idempotent_test",
        "MAMBA_2",
        TrainingStatus::Running,
    );
    state.add_job(job);

    let addr = start_mock_server(state.clone()).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request1 = StopTrainingRequest {
        job_id: "idempotent_test".to_string(),
        reason: "First stop".to_string(),
    };

    let response1 = client.stop_training(request1).await;
    assert!(response1.is_ok());

    // Stop again - should be idempotent
    let request2 = StopTrainingRequest {
        job_id: "idempotent_test".to_string(),
        reason: "Second stop".to_string(),
    };

    let response2 = client.stop_training(request2).await;
    assert!(response2.is_ok(), "Stop should be idempotent");
}

/// TEST 7: Stop with custom reason
#[tokio::test]
async fn test_stop_with_reason() {
    let state = MockTrainingState::default();

    let job = test_fixtures::create_test_job(
        "reason_test",
        "TFT",
        TrainingStatus::Running,
    );
    state.add_job(job);

    let addr = start_mock_server(state).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request = StopTrainingRequest {
        job_id: "reason_test".to_string(),
        reason: "Low validation accuracy".to_string(),
    };

    let response = client.stop_training(request).await;
    assert!(response.is_ok());

    let result = response.unwrap().into_inner();
    assert!(result.success);
    assert!(result.message.contains("stopped"));
}

/// TEST 8: Concurrent stop requests
#[tokio::test]
async fn test_concurrent_stops() {
    let state = MockTrainingState::default();

    // Create multiple jobs
    for i in 0..10 {
        let job = test_fixtures::create_test_job(
            &format!("concurrent_stop_{}", i),
            "DQN",
            TrainingStatus::Running,
        );
        state.add_job(job);
    }

    let addr = start_mock_server(state.clone()).await;

    // Stop all concurrently
    let mut handles = vec![];
    for i in 0..10 {
        let addr_clone = addr;
        let handle = tokio::spawn(async move {
            let channel = Channel::from_shared(format!("http://{}", addr_clone))
                .unwrap()
                .connect()
                .await
                .unwrap();

            let mut client = MlTrainingServiceClient::new(channel);

            let request = StopTrainingRequest {
                job_id: format!("concurrent_stop_{}", i),
                reason: "Concurrent stop test".to_string(),
            };

            client.stop_training(request).await
        });
        handles.push(handle);
    }

    // Wait for all stops
    for handle in handles {
        let result = handle.await.unwrap();
        assert!(result.is_ok());
    }

    // Verify all stopped
    for i in 0..10 {
        let job = state.get_job(&format!("concurrent_stop_{}", i)).unwrap();
        assert_eq!(job.status, TrainingStatus::Stopped as i32);
    }
}

/// TEST 9: Stop response message validation
#[tokio::test]
async fn test_stop_response_message() {
    let state = MockTrainingState::default();

    let job = test_fixtures::create_test_job(
        "message_test",
        "PPO",
        TrainingStatus::Running,
    );
    state.add_job(job);

    let addr = start_mock_server(state).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request = StopTrainingRequest {
        job_id: "message_test".to_string(),
        reason: "Test".to_string(),
    };

    let response = client.stop_training(request).await.unwrap();
    let result = response.into_inner();

    assert!(result.success);
    assert!(!result.message.is_empty());
    assert!(result.message.contains("message_test"));
}

/// TEST 10: Stop job cleanup verification
#[tokio::test]
async fn test_stop_cleanup() {
    let state = MockTrainingState::default();

    let job = test_fixtures::create_test_job(
        "cleanup_test",
        "TFT",
        TrainingStatus::Running,
    );
    state.add_job(job);

    let addr = start_mock_server(state.clone()).await;

    let channel = Channel::from_shared(format!("http://{}", addr))
        .unwrap()
        .connect()
        .await
        .unwrap();

    let mut client = MlTrainingServiceClient::new(channel);

    let request = StopTrainingRequest {
        job_id: "cleanup_test".to_string(),
        reason: "Cleanup test".to_string(),
    };

    let response = client.stop_training(request).await;
    assert!(response.is_ok());

    // Verify job still exists but is stopped
    let job = state.get_job("cleanup_test");
    assert!(job.is_some());
    assert_eq!(job.unwrap().status, TrainingStatus::Stopped as i32);
}