replay system

COBY/tests/test_replay_system.py

@@ -0,0 +1,153 @@
+"""
+Test script for the historical data replay system.
+"""
+
+import asyncio
+import logging
+from datetime import datetime, timedelta
+
+from ..config import Config
+from ..storage.storage_manager import StorageManager
+from ..replay.replay_manager import HistoricalReplayManager
+from ..models.core import ReplayStatus
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+async def test_replay_system():
+    """Test the replay system functionality."""
+    
+    # Initialize components
+    config = Config()
+    storage_manager = StorageManager(config)
+    
+    try:
+        # Initialize storage
+        logger.info("Initializing storage manager...")
+        await storage_manager.initialize()
+        
+        # Initialize replay manager
+        logger.info("Initializing replay manager...")
+        replay_manager = HistoricalReplayManager(storage_manager, config)
+        
+        # Test data range query
+        logger.info("Testing data range query...")
+        data_range = await replay_manager.get_available_data_range("BTCUSDT")
+        if data_range:
+            logger.info(f"Available data range: {data_range['start']} to {data_range['end']}")
+        else:
+            logger.warning("No data available for BTCUSDT")
+            return
+        
+        # Create test replay session
+        logger.info("Creating replay session...")
+        start_time = data_range['start']
+        end_time = start_time + timedelta(minutes=5)  # 5 minute replay
+        
+        session_id = replay_manager.create_replay_session(
+            start_time=start_time,
+            end_time=end_time,
+            speed=10.0,  # 10x speed
+            symbols=["BTCUSDT"],
+            exchanges=["binance"]
+        )
+        
+        logger.info(f"Created session: {session_id}")
+        
+        # Add data callback
+        data_count = 0
+        
+        def data_callback(data):
+            nonlocal data_count
+            data_count += 1
+            if data_count % 100 == 0:
+                logger.info(f"Received {data_count} data points")
+        
+        replay_manager.add_data_callback(session_id, data_callback)
+        
+        # Add status callback
+        def status_callback(session_id, status):
+            logger.info(f"Session {session_id} status changed to: {status.value}")
+        
+        replay_manager.add_status_callback(session_id, status_callback)
+        
+        # Start replay
+        logger.info("Starting replay...")
+        await replay_manager.start_replay(session_id)
+        
+        # Monitor progress
+        while True:
+            session = replay_manager.get_replay_status(session_id)
+            if not session:
+                break
+            
+            if session.status in [ReplayStatus.COMPLETED, ReplayStatus.ERROR, ReplayStatus.STOPPED]:
+                break
+            
+            logger.info(f"Progress: {session.progress:.2%}, Events: {session.events_replayed}")
+            await asyncio.sleep(2)
+        
+        # Final status
+        final_session = replay_manager.get_replay_status(session_id)
+        if final_session:
+            logger.info(f"Final status: {final_session.status.value}")
+            logger.info(f"Total events replayed: {final_session.events_replayed}")
+            logger.info(f"Total data callbacks: {data_count}")
+        
+        # Test session controls
+        logger.info("Testing session controls...")
+        
+        # Create another session for control testing
+        control_session_id = replay_manager.create_replay_session(
+            start_time=start_time,
+            end_time=end_time,
+            speed=1.0,
+            symbols=["BTCUSDT"]
+        )
+        
+        # Start and immediately pause
+        await replay_manager.start_replay(control_session_id)
+        await asyncio.sleep(1)
+        await replay_manager.pause_replay(control_session_id)
+        
+        # Test seek
+        seek_time = start_time + timedelta(minutes=2)
+        success = replay_manager.seek_replay(control_session_id, seek_time)
+        logger.info(f"Seek to {seek_time}: {'success' if success else 'failed'}")
+        
+        # Test speed change
+        success = replay_manager.set_replay_speed(control_session_id, 5.0)
+        logger.info(f"Speed change to 5x: {'success' if success else 'failed'}")
+        
+        # Resume and stop
+        await replay_manager.resume_replay(control_session_id)
+        await asyncio.sleep(2)
+        await replay_manager.stop_replay(control_session_id)
+        
+        # Get statistics
+        stats = replay_manager.get_stats()
+        logger.info(f"Replay manager stats: {stats}")
+        
+        # List all sessions
+        sessions = replay_manager.list_replay_sessions()
+        logger.info(f"Total sessions: {len(sessions)}")
+        
+        # Clean up
+        for session in sessions:
+            replay_manager.delete_replay_session(session.session_id)
+        
+        logger.info("Replay system test completed successfully!")
+        
+    except Exception as e:
+        logger.error(f"Test failed: {e}")
+        raise
+    
+    finally:
+        # Clean up
+        await storage_manager.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(test_replay_system())