foxhunt/services/ml_training_service/tests/data_file_discovery_test.rs

//! Data File Discovery Tests (TDD Red Phase)
//!
//! These tests verify the functionality of discovering Parquet and DBN files
//! for ML training. Tests should FAIL initially until implementation is complete.

use std::fs::{self, File};
use std::io::Write;
use std::path::PathBuf;
use tempfile::TempDir;

// Import the module we're testing (will fail compilation until implemented)
use ml_training_service::data_file_discovery::{DataFileDiscovery, DataFormat};

/// Helper to create a test environment with mock data files
struct TestEnv {
    _temp_dir: TempDir,
    base_path: PathBuf,
}

impl TestEnv {
    fn new() -> Self {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let base_path = temp_dir.path().to_path_buf();
        Self {
            _temp_dir: temp_dir,
            base_path,
        }
    }

    fn create_parquet_file(&self, filename: &str) -> PathBuf {
        let file_path = self.base_path.join(filename);
        let mut file = File::create(&file_path).expect("Failed to create parquet file");
        // Write minimal Parquet magic bytes (just for validation)
        file.write_all(b"PAR1").expect("Failed to write to file");
        file_path
    }

    fn create_dbn_file(&self, filename: &str) -> PathBuf {
        let file_path = self.base_path.join(filename);
        let mut file = File::create(&file_path).expect("Failed to create DBN file");
        // Write minimal DBN header (just for validation)
        file.write_all(b"DBN").expect("Failed to write to file");
        file_path
    }

    fn create_invalid_file(&self, filename: &str) -> PathBuf {
        let file_path = self.base_path.join(filename);
        let mut file = File::create(&file_path).expect("Failed to create invalid file");
        file.write_all(b"INVALID").expect("Failed to write to file");
        file_path
    }
}

#[test]
fn test_discover_existing_parquet_files() {
    let env = TestEnv::new();
    env.create_parquet_file("ES_FUT_180d.parquet");
    env.create_parquet_file("NQ_FUT_180d.parquet");

    let discovery = DataFileDiscovery::new(&env.base_path);

    // Use simple string symbols instead of Asset type
    let assets = vec!["ES_FUT".to_string(), "NQ_FUT".to_string()];
    let result = discovery.discover(&assets);

    assert!(result.is_ok(), "Discovery should succeed for existing Parquet files");
    let files = result.unwrap();
    assert_eq!(files.len(), 2, "Should find 2 Parquet files");

    // Verify ES_FUT file
    let es_file = files.iter().find(|f| f.asset == "ES_FUT").unwrap();
    assert!(matches!(es_file.format, DataFormat::Parquet));
    assert!(es_file.path.exists());
    assert_eq!(es_file.path.file_name().unwrap(), "ES_FUT_180d.parquet");

    // Verify NQ_FUT file
    let nq_file = files.iter().find(|f| f.asset == "NQ_FUT").unwrap();
    assert!(matches!(nq_file.format, DataFormat::Parquet));
    assert!(nq_file.path.exists());
}

#[test]
fn test_discover_existing_dbn_files() {
    let env = TestEnv::new();
    env.create_dbn_file("ES_FUT_180d.dbn");
    env.create_dbn_file("ZN_FUT_90d.dbn");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec!["ES_FUT".to_string(), "ZN_FUT".to_string()];
    let result = discovery.discover(&assets);

    assert!(result.is_ok(), "Discovery should succeed for existing DBN files");
    let files = result.unwrap();
    assert_eq!(files.len(), 2, "Should find 2 DBN files");

    // Verify both files are DBN format
    assert!(files.iter().all(|f| matches!(f.format, DataFormat::DBN)));
}

#[test]
fn test_discover_mixed_formats() {
    let env = TestEnv::new();
    env.create_parquet_file("ES_FUT_180d.parquet");
    env.create_dbn_file("NQ_FUT_180d.dbn");
    env.create_parquet_file("6E_FUT_180d.parquet");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec![
        "ES_FUT".to_string(),
        "NQ_FUT".to_string(),
        "6E_FUT".to_string(),
    ];
    let result = discovery.discover(&assets);

    assert!(result.is_ok(), "Discovery should handle mixed formats");
    let files = result.unwrap();
    assert_eq!(files.len(), 3, "Should find all 3 files");

    // Verify format types
    let es_file = files.iter().find(|f| f.asset == "ES_FUT").unwrap();
    assert!(matches!(es_file.format, DataFormat::Parquet));

    let nq_file = files.iter().find(|f| f.asset == "NQ_FUT").unwrap();
    assert!(matches!(nq_file.format, DataFormat::DBN));

    let e6_file = files.iter().find(|f| f.asset == "6E_FUT").unwrap();
    assert!(matches!(e6_file.format, DataFormat::Parquet));
}

#[test]
fn test_discover_with_custom_base_path() {
    let env = TestEnv::new();
    let custom_dir = env.base_path.join("custom_data");
    fs::create_dir(&custom_dir).expect("Failed to create custom dir");

    let custom_file = custom_dir.join("ES_FUT_180d.parquet");
    let mut file = File::create(&custom_file).expect("Failed to create file");
    file.write_all(b"PAR1").expect("Failed to write");

    let discovery = DataFileDiscovery::new(&custom_dir);

    let assets = vec!["ES_FUT".to_string()];
    let result = discovery.discover(&assets);

    assert!(result.is_ok(), "Discovery should work with custom base path");
    let files = result.unwrap();
    assert_eq!(files.len(), 1);
    assert!(files[0].path.starts_with(&custom_dir));
}

#[test]
fn test_discover_prefers_parquet_over_dbn() {
    let env = TestEnv::new();
    // Create both formats for the same asset
    env.create_parquet_file("ES_FUT_180d.parquet");
    env.create_dbn_file("ES_FUT_180d.dbn");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec!["ES_FUT".to_string()];
    let result = discovery.discover(&assets);

    assert!(result.is_ok());
    let files = result.unwrap();
    assert_eq!(files.len(), 1, "Should only return one file per asset");

    // Verify Parquet is preferred
    assert!(
        matches!(files[0].format, DataFormat::Parquet),
        "Should prefer Parquet over DBN"
    );
    assert_eq!(
        files[0].path.extension().unwrap(),
        "parquet",
        "File extension should be .parquet"
    );
}

#[test]
fn test_discover_missing_file_returns_error() {
    let env = TestEnv::new();
    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec!["NONEXISTENT_FUT".to_string()];
    let result = discovery.discover(&assets);

    assert!(result.is_err(), "Discovery should fail for missing files");
    let err_msg = result.unwrap_err().to_string();
    assert!(
        err_msg.contains("NONEXISTENT_FUT") || err_msg.contains("No data file found"),
        "Error message should mention the missing asset"
    );
}

#[test]
fn test_discover_multiple_assets() {
    let env = TestEnv::new();
    env.create_parquet_file("ES_FUT_180d.parquet");
    env.create_parquet_file("NQ_FUT_180d.parquet");
    env.create_parquet_file("6E_FUT_180d.parquet");
    env.create_parquet_file("ZN_FUT_90d.parquet");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec![
        "ES_FUT".to_string(),
        "NQ_FUT".to_string(),
        "6E_FUT".to_string(),
        "ZN_FUT".to_string(),
    ];
    let result = discovery.discover(&assets);

    assert!(result.is_ok());
    let files = result.unwrap();
    assert_eq!(files.len(), 4, "Should find all 4 assets");

    // Verify all assets are present
    let found_assets: Vec<_> = files.iter().map(|f| f.asset.as_str()).collect();
    assert!(found_assets.contains(&"ES_FUT"));
    assert!(found_assets.contains(&"NQ_FUT"));
    assert!(found_assets.contains(&"6E_FUT"));
    assert!(found_assets.contains(&"ZN_FUT"));
}

#[test]
fn test_discover_validates_file_extension() {
    let env = TestEnv::new();
    env.create_invalid_file("ES_FUT_180d.txt");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec!["ES_FUT".to_string()];
    let result = discovery.discover(&assets);

    assert!(
        result.is_err(),
        "Discovery should reject files with invalid extensions"
    );
}

#[test]
fn test_discover_performance_4_assets() {
    use std::time::Instant;

    let env = TestEnv::new();
    env.create_parquet_file("ES_FUT_180d.parquet");
    env.create_parquet_file("NQ_FUT_180d.parquet");
    env.create_parquet_file("6E_FUT_180d.parquet");
    env.create_parquet_file("ZN_FUT_90d.parquet");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec![
        "ES_FUT".to_string(),
        "NQ_FUT".to_string(),
        "6E_FUT".to_string(),
        "ZN_FUT".to_string(),
    ];

    let start = Instant::now();
    let result = discovery.discover(&assets);
    let elapsed = start.elapsed();

    assert!(result.is_ok());
    assert!(
        elapsed.as_millis() < 5,
        "Discovery should complete in <5ms, took {:?}",
        elapsed
    );
}

#[test]
fn test_find_parquet_returns_correct_path() {
    let env = TestEnv::new();
    env.create_parquet_file("ES_FUT_180d.parquet");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let path = discovery.find_parquet(&"ES_FUT".to_string());

    assert!(path.is_some(), "Should find existing Parquet file");
    let path = path.unwrap();
    assert!(path.exists());
    assert_eq!(path.extension().unwrap(), "parquet");
    assert!(path.file_name().unwrap().to_str().expect("INVARIANT: Path should be valid UTF-8").contains("ES_FUT"));
}

#[test]
fn test_find_parquet_returns_none_for_missing() {
    let env = TestEnv::new();
    let discovery = DataFileDiscovery::new(&env.base_path);

    let path = discovery.find_parquet(&"MISSING_FUT".to_string());

    assert!(path.is_none(), "Should return None for missing Parquet file");
}

#[test]
fn test_find_dbn_returns_correct_path() {
    let env = TestEnv::new();
    env.create_dbn_file("ES_FUT_180d.dbn");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let path = discovery.find_dbn(&"ES_FUT".to_string());

    assert!(path.is_some(), "Should find existing DBN file");
    let path = path.unwrap();
    assert!(path.exists());
    assert_eq!(path.extension().unwrap(), "dbn");
}

#[test]
fn test_find_dbn_returns_none_for_missing() {
    let env = TestEnv::new();
    let discovery = DataFileDiscovery::new(&env.base_path);

    let path = discovery.find_dbn(&"MISSING_FUT".to_string());

    assert!(path.is_none(), "Should return None for missing DBN file");
}

#[test]
fn test_discover_handles_small_parquet_suffix() {
    let env = TestEnv::new();
    env.create_parquet_file("ES_FUT_small.parquet");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec!["ES_FUT".to_string()];
    let result = discovery.discover(&assets);

    // Should find the _small variant
    assert!(result.is_ok());
    let files = result.unwrap();
    assert_eq!(files.len(), 1);
}

#[test]
fn test_discover_handles_various_dbn_patterns() {
    let env = TestEnv::new();
    // Different DBN naming patterns used in the codebase
    env.create_dbn_file("ES_FUT_180d.dbn");
    env.create_dbn_file("NQ_FUT_180d_uncompressed.dbn");
    env.create_dbn_file("ZN_FUT_90d.dbn");

    let discovery = DataFileDiscovery::new(&env.base_path);

    let assets = vec![
        "ES_FUT".to_string(),
        "NQ_FUT".to_string(),
        "ZN_FUT".to_string(),
    ];
    let result = discovery.discover(&assets);

    assert!(result.is_ok());
    let files = result.unwrap();
    assert_eq!(files.len(), 3, "Should handle various DBN naming patterns");
}