gogo2/COBY/tests/test_timescale_manager.py

"""
Tests for TimescaleDB storage manager.
"""

import pytest
import asyncio
from datetime import datetime, timezone
from ..storage.timescale_manager import TimescaleManager
from ..models.core import OrderBookSnapshot, TradeEvent, PriceLevel
from ..config import config


@pytest.fixture
async def storage_manager():
    """Create and initialize storage manager for testing"""
    manager = TimescaleManager()
    await manager.initialize()
    yield manager
    await manager.close()


@pytest.fixture
def sample_orderbook():
    """Create sample order book for testing"""
    return OrderBookSnapshot(
        symbol="BTCUSDT",
        exchange="binance",
        timestamp=datetime.now(timezone.utc),
        bids=[
            PriceLevel(price=50000.0, size=1.5, count=3),
            PriceLevel(price=49999.0, size=2.0, count=5)
        ],
        asks=[
            PriceLevel(price=50001.0, size=1.0, count=2),
            PriceLevel(price=50002.0, size=1.5, count=4)
        ],
        sequence_id=12345
    )


@pytest.fixture
def sample_trade():
    """Create sample trade for testing"""
    return TradeEvent(
        symbol="BTCUSDT",
        exchange="binance",
        timestamp=datetime.now(timezone.utc),
        price=50000.5,
        size=0.1,
        side="buy",
        trade_id="test_trade_123"
    )


class TestTimescaleManager:
    """Test cases for TimescaleManager"""
    
    @pytest.mark.asyncio
    async def test_health_check(self, storage_manager):
        """Test storage health check"""
        is_healthy = await storage_manager.health_check()
        assert is_healthy is True
    
    @pytest.mark.asyncio
    async def test_store_orderbook(self, storage_manager, sample_orderbook):
        """Test storing order book snapshot"""
        result = await storage_manager.store_orderbook(sample_orderbook)
        assert result is True
    
    @pytest.mark.asyncio
    async def test_store_trade(self, storage_manager, sample_trade):
        """Test storing trade event"""
        result = await storage_manager.store_trade(sample_trade)
        assert result is True
    
    @pytest.mark.asyncio
    async def test_get_latest_orderbook(self, storage_manager, sample_orderbook):
        """Test retrieving latest order book"""
        # Store the order book first
        await storage_manager.store_orderbook(sample_orderbook)
        
        # Retrieve it
        retrieved = await storage_manager.get_latest_orderbook(
            sample_orderbook.symbol, 
            sample_orderbook.exchange
        )
        
        assert retrieved is not None
        assert retrieved.symbol == sample_orderbook.symbol
        assert retrieved.exchange == sample_orderbook.exchange
        assert len(retrieved.bids) == len(sample_orderbook.bids)
        assert len(retrieved.asks) == len(sample_orderbook.asks)
    
    @pytest.mark.asyncio
    async def test_batch_store_orderbooks(self, storage_manager):
        """Test batch storing order books"""
        orderbooks = []
        for i in range(5):
            orderbook = OrderBookSnapshot(
                symbol="ETHUSDT",
                exchange="binance",
                timestamp=datetime.now(timezone.utc),
                bids=[PriceLevel(price=3000.0 + i, size=1.0)],
                asks=[PriceLevel(price=3001.0 + i, size=1.0)],
                sequence_id=i
            )
            orderbooks.append(orderbook)
        
        result = await storage_manager.batch_store_orderbooks(orderbooks)
        assert result == 5
    
    @pytest.mark.asyncio
    async def test_batch_store_trades(self, storage_manager):
        """Test batch storing trades"""
        trades = []
        for i in range(5):
            trade = TradeEvent(
                symbol="ETHUSDT",
                exchange="binance",
                timestamp=datetime.now(timezone.utc),
                price=3000.0 + i,
                size=0.1,
                side="buy" if i % 2 == 0 else "sell",
                trade_id=f"test_trade_{i}"
            )
            trades.append(trade)
        
        result = await storage_manager.batch_store_trades(trades)
        assert result == 5
    
    @pytest.mark.asyncio
    async def test_get_storage_stats(self, storage_manager):
        """Test getting storage statistics"""
        stats = await storage_manager.get_storage_stats()
        
        assert isinstance(stats, dict)
        assert 'table_sizes' in stats
        assert 'record_counts' in stats
        assert 'connection_pool' in stats
    
    @pytest.mark.asyncio
    async def test_historical_data_retrieval(self, storage_manager, sample_orderbook, sample_trade):
        """Test retrieving historical data"""
        # Store some data first
        await storage_manager.store_orderbook(sample_orderbook)
        await storage_manager.store_trade(sample_trade)
        
        # Define time range
        start_time = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
        end_time = datetime.now(timezone.utc).replace(hour=23, minute=59, second=59, microsecond=999999)
        
        # Retrieve historical order books
        orderbooks = await storage_manager.get_historical_orderbooks(
            sample_orderbook.symbol,
            sample_orderbook.exchange,
            start_time,
            end_time,
            limit=10
        )
        
        assert isinstance(orderbooks, list)
        
        # Retrieve historical trades
        trades = await storage_manager.get_historical_trades(
            sample_trade.symbol,
            sample_trade.exchange,
            start_time,
            end_time,
            limit=10
        )
        
        assert isinstance(trades, list)


if __name__ == "__main__":
    # Run a simple test
    async def simple_test():
        manager = TimescaleManager()
        await manager.initialize()
        
        # Test health check
        is_healthy = await manager.health_check()
        print(f"Health check: {'PASSED' if is_healthy else 'FAILED'}")
        
        # Test storage stats
        stats = await manager.get_storage_stats()
        print(f"Storage stats: {len(stats)} categories")
        
        await manager.close()
        print("Simple test completed")
    
    asyncio.run(simple_test())