foxhunt/ml/examples/test_dbn_loading.rs

//! Test DBN Loading for DQN Training
//!
//! Simple test to verify DBN files can be loaded and parsed successfully.

use anyhow::Result;
use data::providers::databento::dbn_parser::{DbnParser, ProcessedMessage};
use std::collections::HashMap;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

#[tokio::main]
async fn main() -> Result<()> {
    // Setup logging
    let subscriber = FmtSubscriber::builder()
        .with_max_level(tracing::Level::INFO)
        .finish();
    tracing::subscriber::set_global_default(subscriber)?;

    info!("🚀 Testing DBN file loading");

    // Find DBN files
    let data_dir = "test_data/real/databento/ml_training_small";
    let dbn_files: Vec<_> = std::fs::read_dir(data_dir)?
        .filter_map(|entry| entry.ok())
        .filter(|entry| entry.path().extension().and_then(|s| s.to_str()) == Some("dbn"))
        .collect();

    info!("Found {} DBN files in {}", dbn_files.len(), data_dir);

    if dbn_files.is_empty() {
        return Err(anyhow::anyhow!("No DBN files found"));
    }

    // Create parser
    let parser = DbnParser::new()?;

    // Configure symbol map
    let mut symbol_map = HashMap::new();
    symbol_map.insert(0, "6E.FUT".to_string());
    symbol_map.insert(1, "6E.FUT".to_string());
    parser.update_symbol_map(symbol_map);

    // Configure price scales
    let mut price_scales = HashMap::new();
    price_scales.insert(0, 4);
    price_scales.insert(1, 4);
    parser.update_price_scales(price_scales);

    let mut total_messages = 0;
    let mut total_ohlcv = 0;

    // Load first file only for quick test
    let file_path = dbn_files[0].path();
    info!("Loading: {}", file_path.display());

    let dbn_bytes = std::fs::read(&file_path)?;
    info!("File size: {} bytes", dbn_bytes.len());

    let messages = parser.parse_batch(&dbn_bytes)?;
    info!("Parsed {} messages", messages.len());

    total_messages += messages.len();

    // Count OHLCV messages
    for msg in &messages {
        if let ProcessedMessage::Ohlcv {
            symbol,
            open,
            high,
            low,
            close,
            volume,
            ..
        } = msg
        {
            total_ohlcv += 1;

            // Print first OHLCV bar as sample
            if total_ohlcv == 1 {
                info!("\n📊 Sample OHLCV Bar:");
                info!("  Symbol: {}", symbol);
                info!("  Open:   {:.4}", open.to_f64());
                info!("  High:   {:.4}", high.to_f64());
                info!("  Low:    {:.4}", low.to_f64());
                info!("  Close:  {:.4}", close.to_f64());
                info!("  Volume: {}", volume);
            }
        }
    }

    info!("\n✅ Test Results:");
    info!("  Total messages: {}", total_messages);
    info!("  OHLCV bars: {}", total_ohlcv);
    info!("  Other messages: {}", total_messages - total_ohlcv);

    if total_ohlcv == 0 {
        warn!("⚠️  No OHLCV messages found - training data would be empty!");
        return Err(anyhow::anyhow!("No OHLCV data in DBN file"));
    }

    info!("🎉 DBN loading test passed!");

    Ok(())
}