foxhunt/crates/ml/tests/test_streaming_loader.rs

#![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<f64> {
    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)
}