TCPDashboard/tests/test_real_storage.py

#!/usr/bin/env python3
"""
Test script for real database storage.

This script tests the OKX data collection system with actual database storage
to verify that raw trades and completed candles are being properly stored.
"""

import asyncio
import signal
import sys
import time
from datetime import datetime, timezone

from data.exchanges.okx import OKXCollector
from data.base_collector import DataType
from database.connection import DatabaseConnection
from utils.logger import get_logger

# Global test state
test_state = {
    'running': True,
    'collectors': []
}

def signal_handler(signum, frame):
    """Handle shutdown signals."""
    print(f"\n📡 Received signal {signum}, shutting down collectors...")
    test_state['running'] = False

# Register signal handlers
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)


async def check_database_connection():
    """Check if database connection is available."""
    try:
        db_manager = DatabaseConnection()
        # Test connection
        with db_manager.get_session() as session:
            session.execute("SELECT 1")
        print("✅ Database connection successful")
        return True
    except Exception as e:
        print(f"❌ Database connection failed: {e}")
        print("   Make sure your database is running and configured correctly")
        return False


async def count_stored_data():
    """Count raw trades and candles in database."""
    try:
        db_manager = DatabaseConnection()
        with db_manager.get_session() as session:
            # Count raw trades
            raw_count = session.execute("SELECT COUNT(*) FROM raw_trades WHERE exchange = 'okx'").scalar()
            
            # Count market data candles
            candle_count = session.execute("SELECT COUNT(*) FROM market_data WHERE exchange = 'okx'").scalar()
            
            print(f"📊 Database counts: Raw trades: {raw_count}, Candles: {candle_count}")
            return raw_count, candle_count
    except Exception as e:
        print(f"❌ Error counting database records: {e}")
        return 0, 0


async def test_real_storage(symbol: str = "BTC-USDT", duration: int = 60):
    """Test real database storage for specified duration."""
    logger = get_logger("real_storage_test")
    logger.info(f"🗄️  Testing REAL database storage for {symbol} for {duration} seconds")
    
    # Check database connection first
    if not await check_database_connection():
        logger.error("Cannot proceed without database connection")
        return False
    
    # Get initial counts
    initial_raw, initial_candles = await count_stored_data()
    
    # Create collector with real database storage
    collector = OKXCollector(
        symbol=symbol,
        data_types=[DataType.TRADE, DataType.ORDERBOOK, DataType.TICKER],
        store_raw_data=True
    )
    
    test_state['collectors'].append(collector)
    
    try:
        # Connect and start collection
        logger.info(f"Connecting to OKX for {symbol}...")
        if not await collector.connect():
            logger.error(f"Failed to connect collector for {symbol}")
            return False
        
        if not await collector.subscribe_to_data([symbol], collector.data_types):
            logger.error(f"Failed to subscribe to data for {symbol}")
            return False
        
        if not await collector.start():
            logger.error(f"Failed to start collector for {symbol}")
            return False
        
        logger.info(f"✅ Successfully started real storage test for {symbol}")
        
        # Monitor for specified duration
        start_time = time.time()
        next_check = start_time + 10  # Check every 10 seconds
        
        while time.time() - start_time < duration and test_state['running']:
            await asyncio.sleep(1)
            
            if time.time() >= next_check:
                # Get and log statistics
                stats = collector.get_status()
                logger.info(f"[{symbol}] Stats: "
                           f"Messages: {stats['processing_stats']['messages_received']}, "
                           f"Trades: {stats['processing_stats']['trades_processed']}, "
                           f"Candles: {stats['processing_stats']['candles_processed']}")
                
                # Check database counts
                current_raw, current_candles = await count_stored_data()
                new_raw = current_raw - initial_raw
                new_candles = current_candles - initial_candles
                logger.info(f"[{symbol}] NEW storage: Raw trades: +{new_raw}, Candles: +{new_candles}")
                
                next_check += 10
        
        # Final counts
        final_raw, final_candles = await count_stored_data()
        total_new_raw = final_raw - initial_raw
        total_new_candles = final_candles - initial_candles
        
        logger.info(f"🏁 FINAL RESULTS for {symbol}:")
        logger.info(f"   📈 Raw trades stored: {total_new_raw}")
        logger.info(f"   🕯️  Candles stored: {total_new_candles}")
        
        # Stop collector
        await collector.unsubscribe_from_data([symbol], collector.data_types)
        await collector.stop()
        await collector.disconnect()
        
        logger.info(f"✅ Completed real storage test for {symbol}")
        
        # Return success if we stored some data
        return total_new_raw > 0
        
    except Exception as e:
        logger.error(f"❌ Error in real storage test for {symbol}: {e}")
        return False


async def main():
    """Main test function."""
    print("🗄️  OKX Real Database Storage Test")
    print("=" * 50)
    
    logger = get_logger("main")
    
    try:
        # Test with real database storage
        success = await test_real_storage("BTC-USDT", 60)
        
        if success:
            print("✅ Real storage test completed successfully!")
            print("   Check your database tables:")
            print("   - raw_trades table should have new OKX trade data")
            print("   - market_data table should have new OKX candles")
        else:
            print("❌ Real storage test failed")
            sys.exit(1)
        
    except Exception as e:
        logger.error(f"Test failed: {e}")
        sys.exit(1)
    
    print("Test completed")


if __name__ == "__main__":
    asyncio.run(main())