view/backend/tests/scheduler_test.rs

use backend::{
    Scheduler, DatabaseHandler, SchedulerCommand,
    create_command_channel,
};
use std::sync::Arc;
use tokio::sync::oneshot;
use tokio::time::{timeout, Duration};
use dotenv::dotenv;
use log::info;
use std::env;

fn get_database_url() -> String {
    dotenv().ok();
    env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgres://quant:2Nr!Ya&oVvY5pp@172.18.0.2:5432/crypto_database".to_string()
    })
}

#[tokio::test]
async fn test_scheduler_startup_and_shutdown() {

    dotenv().ok();
    env_logger::try_init().ok();
    
    let db_url = get_database_url();
    let db = Arc::new(DatabaseHandler::new(&db_url).await.expect("Failed to connect to database"));
    
    let (command_tx, command_rx) = create_command_channel();
    let mut scheduler = Scheduler::new(db.clone(), command_rx, None);
    
    // Start scheduler in background
    let scheduler_handle = tokio::spawn(async move {
        scheduler.run().await;
    });
    
    // Give it time to start
    tokio::time::sleep(Duration::from_secs(2)).await;
    
    // Send shutdown command
    command_tx.send(SchedulerCommand::Shutdown).await.expect("Failed to send shutdown");
    
    // Wait for scheduler to stop
    let result = timeout(Duration::from_secs(10), scheduler_handle).await;
    assert!(result.is_ok(), "Scheduler did not shutdown in time");
    
    info!("Scheduler startup and shutdown test passed");
}

#[tokio::test]
async fn test_get_status() {

    dotenv().ok();
    env_logger::try_init().ok();
    
    let db_url = get_database_url();
    let db = Arc::new(DatabaseHandler::new(&db_url).await.expect("Failed to connect to database"));
    
    let (command_tx, command_rx) = create_command_channel();
    let mut scheduler = Scheduler::new(db.clone(), command_rx, None);
    
    // Start scheduler
    let scheduler_handle = tokio::spawn(async move {
        scheduler.run().await;
    });
    
    // Give it time to start
    tokio::time::sleep(Duration::from_secs(2)).await;
    
    // Request status
    let (reply_tx, reply_rx) = oneshot::channel();
    command_tx.send(SchedulerCommand::GetStatus { reply: reply_tx }).await.expect("Failed to send GetStatus");
    
    let status = timeout(Duration::from_secs(5), reply_rx).await
        .expect("Timeout waiting for status")
        .expect("Failed to receive status");
    
    info!("Scheduler status: {:?}", status);
    assert!(status.is_running || status.active_symbols.is_empty(), "Scheduler should be running or have no symbols");
    
    // Shutdown
    command_tx.send(SchedulerCommand::Shutdown).await.expect("Failed to send shutdown");
    let _ = timeout(Duration::from_secs(10), scheduler_handle).await;
    
    info!("GetStatus test passed");
}

#[tokio::test]
async fn test_add_and_remove_symbol() {

    dotenv().ok();
    env_logger::try_init().ok();
    
    let db_url = get_database_url();
    let db = Arc::new(DatabaseHandler::new(&db_url).await.expect("Failed to connect to database"));
    
    let test_symbol = "testscheduler";
    
    // Clean up first - remove test symbol if exists
    let _ = db.remove_symbol(test_symbol).await;
    
    let (command_tx, command_rx) = create_command_channel();
    let mut scheduler = Scheduler::new(db.clone(), command_rx, None);
    
    // Start scheduler
    let scheduler_handle = tokio::spawn(async move {
        scheduler.run().await;
    });
    
    // Give it time to start
    tokio::time::sleep(Duration::from_secs(2)).await;
    
    // Add symbol (with recent backfill_from to avoid long backfill)
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64;
    let recent_time = now - (5 * 60 * 1000); // 5 minutes ago
    
    info!("Adding test symbol: {}", test_symbol);
    command_tx.send(SchedulerCommand::AddSymbol { 
        symbol: test_symbol.to_string(),
        backfill_from: Some(recent_time),
    }).await.expect("Failed to send AddSymbol");
    
    // Wait for processing
    tokio::time::sleep(Duration::from_secs(5)).await;
    
    // Check status
    let (reply_tx, reply_rx) = oneshot::channel();
    command_tx.send(SchedulerCommand::GetStatus { reply: reply_tx }).await.expect("Failed to send GetStatus");
    
    let status = timeout(Duration::from_secs(5), reply_rx).await
        .expect("Timeout")
        .expect("Failed to receive");
    
    info!("Status after add: {:?}", status);
    
    // Verify symbol was added
    let is_tracked = db.is_symbol_tracked(test_symbol).await.expect("Failed to check tracking");
    assert!(is_tracked, "Symbol should be tracked after AddSymbol");
    
    // Remove symbol
    info!("Removing test symbol: {}", test_symbol);
    command_tx.send(SchedulerCommand::RemoveSymbol { 
        symbol: test_symbol.to_string(),
    }).await.expect("Failed to send RemoveSymbol");
    
    // Wait for processing
    tokio::time::sleep(Duration::from_secs(3)).await;
    
    // Verify symbol was removed
    let is_tracked = db.is_symbol_tracked(test_symbol).await.expect("Failed to check tracking");
    assert!(!is_tracked, "Symbol should not be tracked after RemoveSymbol");
    
    // Shutdown
    command_tx.send(SchedulerCommand::Shutdown).await.expect("Failed to send shutdown");
    let _ = timeout(Duration::from_secs(10), scheduler_handle).await;
    
    info!("Add and remove symbol test passed");
}

#[tokio::test]
async fn test_restart_collector() {

    dotenv().ok();
    env_logger::try_init().ok();
    
    let db_url = get_database_url();
    let db = Arc::new(DatabaseHandler::new(&db_url).await.expect("Failed to connect to database"));
    
    let (command_tx, command_rx) = create_command_channel();
    let mut scheduler = Scheduler::new(db.clone(), command_rx, None);
    
    // Start scheduler
    let scheduler_handle = tokio::spawn(async move {
        scheduler.run().await;
    });
    
    // Give it time to start
    tokio::time::sleep(Duration::from_secs(2)).await;
    
    // Get initial status
    let (reply_tx, reply_rx) = oneshot::channel();
    command_tx.send(SchedulerCommand::GetStatus { reply: reply_tx }).await.unwrap();
    let status_before = reply_rx.await.unwrap();
    info!("Status before restart: {:?}", status_before);
    
    // Send restart command
    info!("Sending RestartCollector command");
    command_tx.send(SchedulerCommand::RestartCollector).await.expect("Failed to send RestartCollector");
    
    // Wait for restart
    tokio::time::sleep(Duration::from_secs(5)).await;
    
    // Get status after restart
    let (reply_tx, reply_rx) = oneshot::channel();
    command_tx.send(SchedulerCommand::GetStatus { reply: reply_tx }).await.unwrap();
    let status_after = reply_rx.await.unwrap();
    info!("Status after restart: {:?}", status_after);
    
    // Shutdown
    command_tx.send(SchedulerCommand::Shutdown).await.expect("Failed to send shutdown");
    let _ = timeout(Duration::from_secs(10), scheduler_handle).await;
    
    info!("Restart collector test passed");
}

#[tokio::test]
async fn test_multiple_commands() {

    dotenv().ok();
    env_logger::try_init().ok();
    
    let db_url = get_database_url();
    let db = Arc::new(DatabaseHandler::new(&db_url).await.expect("Failed to connect to database"));
    
    let (command_tx, command_rx) = create_command_channel();
    let mut scheduler = Scheduler::new(db.clone(), command_rx, None);
    
    // Start scheduler
    let scheduler_handle = tokio::spawn(async move {
        scheduler.run().await;
    });
    
    tokio::time::sleep(Duration::from_secs(2)).await;
    
    // Send multiple status requests rapidly
    for i in 0..5 {
        let (reply_tx, reply_rx) = oneshot::channel();
        command_tx.send(SchedulerCommand::GetStatus { reply: reply_tx }).await.unwrap();
        let status = reply_rx.await.unwrap();
        info!("Status request {}: {:?}", i, status);
    }
    
    // Shutdown
    command_tx.send(SchedulerCommand::Shutdown).await.expect("Failed to send shutdown");
    let _ = timeout(Duration::from_secs(10), scheduler_handle).await;
    
    info!("Multiple commands test passed");
}

#[tokio::test]
async fn test_real_btcusdt_collection() {

    dotenv().ok();
    env_logger::try_init().ok();
    
    let db_url = get_database_url();
    let db = Arc::new(DatabaseHandler::new(&db_url).await.expect("Failed to connect to database"));
    
    let symbol = "btcusdt";
    
    // Clean up - remove from tracking first
    let _ = db.remove_symbol(symbol).await;
    
    let (command_tx, command_rx) = create_command_channel();
    let mut scheduler = Scheduler::new(db.clone(), command_rx, None);
    
    // Start scheduler
    let scheduler_handle = tokio::spawn(async move {
        scheduler.run().await;
    });
    
    tokio::time::sleep(Duration::from_secs(2)).await;
    
    // Calculate start time: 5 minutes ago
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64;
    let five_min_ago = now - (5 * 60 * 1000);
    
    info!("Adding BTCUSDT, backfill from {} (5 min ago)", five_min_ago);
    
    // Add BTCUSDT with backfill from 5 minutes ago
    command_tx.send(SchedulerCommand::AddSymbol { 
        symbol: symbol.to_string(),
        backfill_from: Some(five_min_ago),
    }).await.expect("Failed to send AddSymbol");
    
    // Wait for backfill to complete
    info!("Waiting for backfill...");
    tokio::time::sleep(Duration::from_secs(10)).await;
    
    // Check latest timestamp in database
    let latest_ts = db.get_latest_timestamp(symbol).await.expect("Failed to get timestamp");
    info!("Latest BTCUSDT timestamp in DB: {:?}", latest_ts);
    
    assert!(latest_ts.is_some(), "Should have BTCUSDT data in database");
    
    // Wait for 2 more candles (about 2 minutes + buffer)
    info!("Waiting for 2 live candles (~2.5 minutes)...");
    tokio::time::sleep(Duration::from_secs(150)).await;
    
    // Check new latest timestamp
    let new_latest_ts = db.get_latest_timestamp(symbol).await.expect("Failed to get timestamp");
    info!("New latest BTCUSDT timestamp: {:?}", new_latest_ts);
    
    assert!(new_latest_ts.unwrap() > latest_ts.unwrap(), "Should have received new candles");
    
    // Get status
    let (reply_tx, reply_rx) = oneshot::channel();
    command_tx.send(SchedulerCommand::GetStatus { reply: reply_tx }).await.unwrap();
    let status = reply_rx.await.unwrap();
    info!("Final status: {:?}", status);
    
    assert!(status.active_symbols.contains(&symbol.to_string()), "BTCUSDT should be in active symbols");
    
    // Clean up - remove symbol
    command_tx.send(SchedulerCommand::RemoveSymbol { 
        symbol: symbol.to_string(),
    }).await.expect("Failed to remove symbol");
    
    tokio::time::sleep(Duration::from_secs(2)).await;
    
    // Shutdown
    command_tx.send(SchedulerCommand::Shutdown).await.expect("Failed to send shutdown");
    let _ = timeout(Duration::from_secs(10), scheduler_handle).await;
    
    info!("Real BTCUSDT collection test passed!");
    info!("Check database: SELECT * FROM klines_1m WHERE symbol = 'BTCUSDT' ORDER BY timestamp DESC LIMIT 10;");
}