rustytorch/crates/training/rtx-automeasure/tests/lib_tests.rs

use rtx_automeasure::{AutoMLAgent, AutoMLConfig, AutoMLPipeline, OptimizationObjective, TaskType};
use rtx_tensor::{DType, Device, Tensor};

#[tokio::test]
async fn test_automl_config_creation() {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(3600) // 1 hour
        .with_memory_budget(8 * 1024 * 1024 * 1024) // 8GB
        .with_objective(OptimizationObjective::Accuracy)
        .with_cv_folds(5);

    assert_eq!(config.task_type, TaskType::Classification);
    assert_eq!(config.time_budget_seconds, 3600);
    assert_eq!(config.memory_budget_bytes, 8 * 1024 * 1024 * 1024);
    assert_eq!(config.cv_folds, 5);
}

#[tokio::test]
async fn test_automl_agent_creation() {
    let config = AutoMLConfig::default();
    let agent = AutoMLAgent::new(config);

    assert!(agent.is_ok());
    let agent = agent.unwrap();
    assert!(!agent.get_id().is_empty());
}

#[tokio::test]
async fn test_automl_agent_fit() {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(60); // 1 minute for test

    let mut agent = AutoMLAgent::new(config).unwrap();

    // Create sample data
    let device = Device::cpu();
    let x_train = Tensor::randn(&[100, 10], &device).unwrap();
    let y_train = Tensor::zeros_typed([100], DType::I64, &device).unwrap();

    let result = agent.fit(&x_train, &y_train).await;
    assert!(result.is_ok());

    let pipeline = result.unwrap();
    assert!(!pipeline.get_models().is_empty());
}

#[tokio::test]
async fn test_automl_agent_predict() {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Regression)
        .with_time_budget(60);

    let mut agent = AutoMLAgent::new(config).unwrap();

    // Create sample data
    let device = Device::cpu();
    let x_train = Tensor::randn(&[50, 5], &device).unwrap();
    let y_train = Tensor::randn(&[50], &device).unwrap();
    let x_test = Tensor::randn(&[20, 5], &device).unwrap();

    let pipeline = agent.fit(&x_train, &y_train).await.unwrap();
    let predictions = agent.predict(&pipeline, &x_test).await;

    assert!(predictions.is_ok());
    let pred_tensor = predictions.unwrap();
    assert_eq!(pred_tensor.shape()[0], 20);
}

#[tokio::test]
async fn test_automl_pipeline_serialization() {
    let config = AutoMLConfig::new().with_task_type(TaskType::Classification);

    let mut agent = AutoMLAgent::new(config).unwrap();

    let device = Device::cpu();
    let x_train = Tensor::randn(&[30, 4], &device).unwrap();
    let y_train = Tensor::zeros_typed([30], DType::I64, &device).unwrap();

    let pipeline = agent.fit(&x_train, &y_train).await.unwrap();

    // Test serialization
    let serialized = pipeline.to_json();
    assert!(serialized.is_ok());

    // Test deserialization
    let json_str = serialized.unwrap();
    let deserialized = AutoMLPipeline::from_json(&json_str);
    assert!(deserialized.is_ok());
}

#[tokio::test]
async fn test_automl_agent_get_leaderboard() {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(30);

    let mut agent = AutoMLAgent::new(config).unwrap();

    let device = Device::cpu();
    let x_train = Tensor::randn(&[40, 3], &device).unwrap();
    let y_train = Tensor::zeros_typed([40], DType::I64, &device).unwrap();

    let _pipeline = agent.fit(&x_train, &y_train).await.unwrap();

    let leaderboard = agent.get_leaderboard();
    assert!(!leaderboard.is_empty());

    // Check that leaderboard is sorted by score
    for i in 1..leaderboard.len() {
        assert!(leaderboard[i - 1].score >= leaderboard[i].score);
    }
}

#[tokio::test]
async fn test_automl_agent_get_feature_importance() {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Regression)
        .with_time_budget(30);

    let mut agent = AutoMLAgent::new(config).unwrap();

    let device = Device::cpu();
    let x_train = Tensor::randn(&[50, 8], &device).unwrap();
    let y_train = Tensor::randn(&[50], &device).unwrap();

    let pipeline = agent.fit(&x_train, &y_train).await.unwrap();

    let importance = agent.get_feature_importance(&pipeline);
    assert!(importance.is_ok());

    let importance_map = importance.unwrap();
    assert_eq!(importance_map.len(), 8); // 8 features
}

#[tokio::test]
async fn test_automl_config_validation() {
    // Test invalid time budget
    let config = AutoMLConfig::new().with_time_budget(0);
    let agent_result = AutoMLAgent::new(config);
    assert!(agent_result.is_err());

    // Test invalid CV folds
    let config = AutoMLConfig::new().with_cv_folds(0);
    let agent_result = AutoMLAgent::new(config);
    assert!(agent_result.is_err());
}

#[tokio::test]
async fn test_automl_agent_progress_tracking() {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(45);

    let mut agent = AutoMLAgent::new(config).unwrap();

    let device = Device::cpu();
    let x_train = Tensor::randn(&[60, 6], &device).unwrap();
    let y_train = Tensor::zeros_typed([60], DType::I64, &device).unwrap();

    // Start fitting
    let _pipeline = agent.fit(&x_train, &y_train).await.unwrap();

    // Check progress after fitting
    let progress = agent.get_progress();

    assert!(progress.elapsed_time_seconds >= 0.0);
    assert!(progress.completion_percentage >= 0.0);
    assert!(progress.completion_percentage <= 100.0);
}