poimen-memory/tests/it_m5_training.rs

use mem_core::{VerlTrainingConfig, TrainingResult, Trajectory};
use mem_llm::VllmConfig;

/// M5.4-M5.6 Integration Tests — vLLM Setup + Training Loop + Gate
/// 
/// Verifies:
/// - vLLM configuration for LoRA
/// - Training hyperparameter validation
/// - Trajectory compatibility with training
/// - Gate criteria (return-over-baseline)

#[test]
fn a1_vllm_config_default() {
    let config = VllmConfig::default();
    
    assert_eq!(config.base_model, "qwen2.5-3b-instruct");
    assert_eq!(config.served_model_name, "memory");
    assert_eq!(config.max_lora_rank, 32);
    assert_eq!(config.max_model_len, 32768);
}

#[test]
fn a2_vllm_adapter_mounting() {
    let mut config = VllmConfig::default();
    config.add_adapter("memory-v1".to_string(), "/mnt/adapters/memory-v1".to_string());
    config.add_adapter("memory-v2".to_string(), "/mnt/adapters/memory-v2".to_string());
    
    assert_eq!(config.lora_modules.len(), 2);
    assert!(config.lora_modules.contains_key("memory-v1"));
    assert!(config.lora_modules.contains_key("memory-v2"));
}

#[test]
fn a3_vllm_container_args() {
    let mut config = VllmConfig::default();
    config.add_adapter("memory-v1".to_string(), "/mnt/adapters/memory-v1".to_string());
    
    let args = config.to_container_args();
    
    // Should include essential flags
    assert!(args.contains(&"python".to_string()));
    assert!(args.contains(&"--enable-lora".to_string()));
    assert!(args.contains(&"--max-lora-rank".to_string()));
    assert!(args.contains(&"32".to_string()));
}

#[test]
fn a4_training_config_default() {
    let config = VerlTrainingConfig::default();
    
    assert_eq!(config.lora_rank, 32);
    assert_eq!(config.train_batch_size, 8);
    assert_eq!(config.num_train_epochs, 3);
    
    // Loss weights should sum to 1.0
    let total = config.trajectory_loss_weight + config.turn_loss_weight;
    assert!((total - 1.0).abs() < 0.01);
}

#[test]
fn a5_training_config_validates() {
    let config = VerlTrainingConfig::default();
    assert!(config.validate().is_ok());
}

#[test]
fn a6_training_config_rejects_invalid_lr() {
    let config = VerlTrainingConfig {
        learning_rate: 1e-9,  // Too low
        ..Default::default()
    };
    assert!(config.validate().is_err());
}

#[test]
fn a7_effective_batch_size() {
    let config = VerlTrainingConfig {
        train_batch_size: 8,
        gradient_accumulation_steps: 4,
        ..Default::default()
    };
    
    assert_eq!(config.effective_batch_size(), 32);
}

#[test]
fn a8_training_scales_to_corpus_size() {
    let small = VerlTrainingConfig::from_corpus("corpus.jsonl", 50, 3);
    let large = VerlTrainingConfig::from_corpus("corpus.jsonl", 2000, 3);
    
    // Large corpus should use bigger batches
    assert!(large.train_batch_size >= small.train_batch_size);
}

#[test]
fn a9_trajectory_compatible_with_training() {
    let mut traj = Trajectory::new("run_001".to_string());
    
    // Add turns with rewards
    for t in 1..=10 {
        let r_update = if t % 2 == 0 { 1 } else { -1 };
        traj.add_turn(t, format!("prompt_{}", t), format!("response_{}", t), r_update, true);
    }
    
    traj.set_exit_reward(5, 5);
    
    // Should serialize for JSONL export
    let json = serde_json::to_string(&traj).expect("Should serialize");
    assert!(json.contains("run_001"));
    
    // Should have correct rewards
    assert_eq!(traj.r_format, 1.0, "All turns parsed");
    assert_eq!(traj.r_exit, 0.0, "Exited at evidence");
}

#[test]
fn a10_training_result_structure() {
    let result = TrainingResult {
        final_loss: 0.45,
        steps_trained: 1000,
        checkpoint_path: "/checkpoints/memory-v1".to_string(),
        epoch: 2,
        timestamp: "2026-08-25T20:00:00Z".to_string(),
    };
    
    assert!(result.final_loss > 0.0);
    assert!(!result.checkpoint_path.is_empty());
    assert_eq!(result.epoch, 2);
}

#[test]
fn a11_gate_criteria_defined() {
    // M5.6 gate checks return-over-baseline
    // Structure for verification:
    struct GateCriteria {
        min_return_improvement: f32,     // Minimum % improvement
        max_training_loss: f32,          // Max acceptable final loss
        min_success_rate: f32,           // Min % of test trajectories passing
    }
    
    let gate = GateCriteria {
        min_return_improvement: 0.1,     // 10% better than baseline
        max_training_loss: 0.5,
        min_success_rate: 0.75,          // 75% of tests should pass
    };
    
    assert!(gate.min_return_improvement > 0.0);
    assert!(gate.max_training_loss > 0.0);
    assert!(gate.min_success_rate > 0.0 && gate.min_success_rate < 1.0);
}

#[test]
fn a12_vllm_endpoint_configuration() {
    let config = VllmConfig {
        endpoint: "http://memory-serving.llm-serving.svc.cluster.local:8000/v1".to_string(),
        api_key: Some("sk-test-key-12345".to_string()),
        ..Default::default()
    };
    
    assert!(config.endpoint.contains("memory"));
    assert!(config.api_key.is_some());
}

#[test]
fn a13_training_hyperparameter_sweep() {
    let learning_rates = vec![1e-5, 5e-5, 1e-4];
    let batch_sizes = vec![4, 8, 16];
    
    let mut configs = Vec::new();
    for lr in learning_rates {
        for bs in &batch_sizes {
            let config = VerlTrainingConfig {
                learning_rate: lr,
                train_batch_size: *bs,
                ..Default::default()
            };
            configs.push(config);
        }
    }
    
    assert_eq!(configs.len(), 9, "3x3 hyperparameter sweep");
    
    // All should validate
    for config in configs {
        assert!(config.validate().is_ok());
    }
}

#[test]
fn a14_checkpoint_management() {
    let checkpoints = vec![
        "/checkpoints/memory-v1-epoch1",
        "/checkpoints/memory-v1-epoch2",
        "/checkpoints/memory-v1-best",
    ];
    
    assert_eq!(checkpoints.len(), 3);
    assert!(checkpoints.iter().all(|p| p.contains("memory")));
}

#[test]
fn a15_m5_completion_status() {
    // Verify all three M5.4-M5.6 phases have structures
    let vllm = VllmConfig::default();
    let training = VerlTrainingConfig::default();
    let result = TrainingResult {
        final_loss: 0.4,
        steps_trained: 500,
        checkpoint_path: "/tmp/checkpoint".to_string(),
        epoch: 1,
        timestamp: "2026-08-25T00:00:00Z".to_string(),
    };
    
    // All required structures present
    assert!(!vllm.base_model.is_empty());
    assert!(training.validate().is_ok());
    assert!(result.final_loss > 0.0);
}