#![allow( clippy::assertions_on_constants, clippy::assertions_on_result_states, clippy::clone_on_copy, clippy::decimal_literal_representation, clippy::doc_markdown, clippy::empty_line_after_doc_comments, clippy::field_reassign_with_default, clippy::get_unwrap, clippy::identity_op, clippy::inconsistent_digit_grouping, clippy::indexing_slicing, clippy::integer_division, clippy::len_zero, clippy::let_underscore_must_use, clippy::manual_div_ceil, clippy::manual_let_else, clippy::manual_range_contains, clippy::modulo_arithmetic, clippy::needless_range_loop, clippy::non_ascii_literal, clippy::redundant_clone, clippy::shadow_reuse, clippy::shadow_same, clippy::shadow_unrelated, clippy::single_match_else, clippy::str_to_string, clippy::string_slice, clippy::tests_outside_test_module, clippy::too_many_lines, clippy::unnecessary_wraps, clippy::unseparated_literal_suffix, clippy::use_debug, clippy::useless_vec, clippy::wildcard_enum_match_arm, clippy::else_if_without_else, clippy::expect_used, clippy::missing_const_for_fn, clippy::similar_names, clippy::type_complexity, clippy::collapsible_else_if, clippy::doc_lazy_continuation, clippy::items_after_test_module, clippy::map_clone, clippy::multiple_unsafe_ops_per_block, clippy::unwrap_or_default, clippy::assign_op_pattern, clippy::needless_borrow, clippy::println_empty_string, clippy::unnecessary_cast, clippy::used_underscore_binding, clippy::create_dir, clippy::implicit_saturating_sub, clippy::exit, clippy::expect_fun_call, clippy::too_many_arguments, clippy::unnecessary_map_or, clippy::unwrap_used, dead_code, unused_imports, unused_variables, clippy::cloned_ref_to_slice_refs, clippy::neg_multiply, clippy::while_let_loop, clippy::bool_assert_comparison, clippy::excessive_precision, clippy::trivially_copy_pass_by_ref, clippy::op_ref, clippy::redundant_closure, clippy::unnecessary_lazy_evaluations, clippy::if_then_some_else_none, clippy::unnecessary_to_owned, clippy::single_component_path_imports, )] //! Integration tests for StreamingDbnLoader //! //! Tests memory-efficient streaming data loading with real DBN files. use anyhow::Result; use ml::data_loaders::{DbnSequenceLoader, StreamingDbnLoader}; use std::path::PathBuf; use tracing::info; /// Test data directory (small dataset with 4 files) const TEST_DATA_DIR: &str = "test_data/real/databento/ml_training_small"; #[tokio::test] async fn test_streaming_loader_creation() -> Result<()> { let loader = StreamingDbnLoader::new(60, 256).await?; info!(?loader, "StreamingDbnLoader created successfully"); Ok(()) } #[tokio::test] async fn test_custom_config() -> Result<()> { let loader = StreamingDbnLoader::with_config(60, 256, 5000, 50).await?; info!(?loader, "Custom config applied"); Ok(()) } #[tokio::test] async fn test_stream_sequences_small_dataset() -> Result<()> { let test_dir = PathBuf::from(TEST_DATA_DIR); if !test_dir.exists() { info!("Test data not found, skipping test"); return Ok(()); } let loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?; let mut stream = loader.stream_sequences(&test_dir, 0.9).await?; let mut total_sequences = 0; let mut batch_count = 0; // Process all batches loop { match stream.next_batch().await? { Some(batch) => { batch_count += 1; total_sequences += batch.len(); // Verify batch contents assert!(!batch.is_empty(), "Batch should not be empty"); for (input, target) in &batch { // Verify tensor shapes assert_eq!(input.dims().len(), 2, "Input should be 2D"); assert_eq!(input.dims()[0], 60, "Sequence length should be 60"); assert_eq!(input.dims()[1], 256, "Feature dim should be 256"); assert_eq!(target.dims().len(), 2, "Target should be 2D"); assert_eq!(target.dims()[0], 1, "Target batch size should be 1"); assert_eq!(target.dims()[1], 256, "Target dim should be 256"); } info!(batch = batch_count, sequences = batch.len(), "Batch processed"); }, None => break, } } info!(total_sequences, batch_count, "Streamed sequences complete"); assert!(total_sequences > 0, "Should load at least some sequences"); assert!(batch_count > 0, "Should have at least one batch"); Ok(()) } #[tokio::test] async fn test_streaming_vs_batch_consistency() -> Result<()> { let test_dir = PathBuf::from(TEST_DATA_DIR); if !test_dir.exists() { info!("Test data not found, skipping test"); return Ok(()); } // Load with batch loader let mut batch_loader = DbnSequenceLoader::with_limits(60, 256, Some(100), 10).await?; let (batch_train, batch_val) = batch_loader.load_sequences(&test_dir, 0.9).await?; let batch_total = batch_train.len() + batch_val.len(); info!(sequences = batch_total, "Batch loader loaded"); // Load with streaming loader (same config) let streaming_loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?; let mut stream = streaming_loader.stream_sequences(&test_dir, 0.9).await?; let mut streaming_total = 0; loop { match stream.next_batch().await? { Some(batch) => streaming_total += batch.len(), None => break, } } info!(sequences = streaming_total, "Streaming loader loaded"); // Should produce similar number of sequences (within 10% due to boundary effects) let diff_ratio = (batch_total as f64 - streaming_total as f64).abs() / batch_total as f64; assert!( diff_ratio < 0.1, "Sequence count should be similar (diff: {:.1}%)", diff_ratio * 100.0 ); info!(diff_pct = diff_ratio * 100.0, "Batch and streaming loaders produce consistent results"); Ok(()) } #[tokio::test] async fn test_memory_efficiency() -> Result<()> { let test_dir = PathBuf::from(TEST_DATA_DIR); if !test_dir.exists() { info!("Test data not found, skipping test"); return Ok(()); } // Get baseline memory let baseline = get_memory_usage_mb()?; info!(baseline_mb = baseline, "Baseline memory"); // Load with streaming let loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?; let mut stream = loader.stream_sequences(&test_dir, 0.9).await?; let mut max_memory = baseline; // Process batches and track peak memory loop { match stream.next_batch().await? { Some(_batch) => { let current = get_memory_usage_mb()?; if current > max_memory { max_memory = current; } }, None => break, } } let peak_memory = max_memory - baseline; info!(peak_mb = peak_memory, "Peak memory delta"); // For small dataset, peak should be < 100MB assert!( peak_memory < 100.0, "Peak memory should be < 100MB for small dataset, got {:.1} MB", peak_memory ); info!(peak_mb = peak_memory, "Memory efficiency verified"); Ok(()) } #[tokio::test] async fn test_train_val_split() -> Result<()> { let test_dir = PathBuf::from(TEST_DATA_DIR); if !test_dir.exists() { info!("Test data not found, skipping test"); return Ok(()); } let loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?; let mut stream = loader.stream_sequences(&test_dir, 0.8).await?; // Count training sequences let mut train_count = 0; loop { match stream.next_batch().await? { Some(batch) => train_count += batch.len(), None => break, } } info!(sequences = train_count, "Training sequences loaded"); // Switch to validation stream.switch_to_validation().await?; // Count validation sequences let mut val_count = 0; loop { match stream.next_batch().await? { Some(batch) => val_count += batch.len(), None => break, } } info!(sequences = val_count, "Validation sequences loaded"); // Verify split ratio is approximately correct (within 20% due to boundary effects) let total = train_count + val_count; let train_ratio = train_count as f64 / total as f64; let split_error = (train_ratio - 0.8).abs(); assert!( split_error < 0.2, "Train/val split should be approximately 80/20, got {:.1}%/{:.1}%", train_ratio * 100.0, (1.0 - train_ratio) * 100.0 ); info!( train_pct = train_ratio * 100.0, val_pct = (1.0 - train_ratio) * 100.0, "Train/val split verified" ); Ok(()) } #[tokio::test] async fn test_different_batch_sizes() -> Result<()> { let test_dir = PathBuf::from(TEST_DATA_DIR); if !test_dir.exists() { info!("Test data not found, skipping test"); return Ok(()); } // Test with different batch sizes let batch_sizes = vec![1000, 5000, 10000]; for batch_size in batch_sizes { let loader = StreamingDbnLoader::with_config(60, 256, batch_size, 10).await?; let mut stream = loader.stream_sequences(&test_dir, 0.9).await?; let mut total = 0; loop { match stream.next_batch().await? { Some(batch) => total += batch.len(), None => break, } } info!(batch_size, total_sequences = total, "Batch size loaded sequences"); assert!( total > 0, "Should load sequences with batch_size={}", batch_size ); } info!("All batch sizes work correctly"); Ok(()) } #[tokio::test] async fn test_different_strides() -> Result<()> { let test_dir = PathBuf::from(TEST_DATA_DIR); if !test_dir.exists() { info!("Test data not found, skipping test"); return Ok(()); } // Test with different strides let strides = vec![1, 10, 50, 100]; for stride in strides { let loader = StreamingDbnLoader::with_config(60, 256, 1000, stride).await?; let mut stream = loader.stream_sequences(&test_dir, 0.9).await?; let mut total = 0; loop { match stream.next_batch().await? { Some(batch) => total += batch.len(), None => break, } } info!(stride, total_sequences = total, "Stride loaded sequences"); assert!(total > 0, "Should load sequences with stride={}", stride); } info!("All strides work correctly"); Ok(()) } /// Get current memory usage in MB fn get_memory_usage_mb() -> Result { let status = std::fs::read_to_string("/proc/self/status")?; for line in status.lines() { if line.starts_with("VmRSS:") { let kb: usize = line .split_whitespace() .nth(1) .and_then(|s| s.parse().ok()) .unwrap_or(0); return Ok(kb as f64 / 1024.0); } } Ok(0.0) }