//! 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"); }