gogo2/COBY/tests/test_kraken_connector.py

"""
Unit tests for Kraken exchange connector.
"""

import asyncio
import pytest
from unittest.mock import Mock, AsyncMock, patch
from datetime import datetime, timezone

from ..connectors.kraken_connector import KrakenConnector
from ..models.core import OrderBookSnapshot, TradeEvent, PriceLevel


class TestKrakenConnector:
    """Test suite for Kraken connector."""
    
    @pytest.fixture
    def connector(self):
        """Create connector instance for testing."""
        return KrakenConnector()
    
    def test_initialization(self, connector):
        """Test connector initializes correctly."""
        assert connector.exchange_name == "kraken"
        assert connector.WEBSOCKET_URL in connector.websocket_url
        assert 'book-25' in connector.message_handlers
        assert 'trade' in connector.message_handlers
        assert connector.system_status == 'unknown'
    
    def test_symbol_normalization(self, connector):
        """Test symbol normalization to Kraken format."""
        # Test standard conversions
        assert connector.normalize_symbol('BTCUSDT') == 'XBT/USD'
        assert connector.normalize_symbol('ETHUSDT') == 'ETH/USD'
        assert connector.normalize_symbol('ADAUSDT') == 'ADA/USD'
        
        # Test generic conversion
        assert connector.normalize_symbol('LINKUSDT') == 'LINK/USD'
        
        # Test already correct format
        assert connector.normalize_symbol('XBT/USD') == 'XBT/USD'
    
    def test_symbol_denormalization(self, connector):
        """Test converting Kraken format back to standard."""
        assert connector._denormalize_symbol('XBT/USD') == 'BTCUSDT'
        assert connector._denormalize_symbol('ETH/USD') == 'ETHUSDT'
        assert connector._denormalize_symbol('ADA/USD') == 'ADAUSDT'
        
        # Test other quote currencies
        assert connector._denormalize_symbol('BTC/EUR') == 'BTCEUR'
    
    def test_message_type_detection(self, connector):
        """Test message type detection."""
        # Test order book message (array format)
        book_message = [123, {'b': [['50000', '1.5']]}, 'book-25', 'XBT/USD']
        assert connector._get_message_type(book_message) == 'book-25'
        
        # Test trade message (array format)
        trade_message = [456, [['50000', '0.1', 1609459200, 'b', 'm', '']], 'trade', 'XBT/USD']
        assert connector._get_message_type(trade_message) == 'trade'
        
        # Test status message (object format)
        status_message = {'event': 'systemStatus', 'status': 'online'}
        assert connector._get_message_type(status_message) == 'systemStatus'
        
        # Test subscription message
        sub_message = {'event': 'subscriptionStatus', 'status': 'subscribed'}
        assert connector._get_message_type(sub_message) == 'subscriptionStatus'
        
        # Test unknown message
        unknown_message = {'data': 'something'}
        assert connector._get_message_type(unknown_message) == 'unknown'
    
    @pytest.mark.asyncio
    async def test_subscription_methods(self, connector):
        """Test subscription and unsubscription methods."""
        # Mock the _send_message method
        connector._send_message = AsyncMock(return_value=True)
        
        # Test order book subscription
        await connector.subscribe_orderbook('BTCUSDT')
        
        # Verify subscription was tracked
        assert 'BTCUSDT' in connector.subscriptions
        assert 'orderbook' in connector.subscriptions['BTCUSDT']
        
        # Verify correct message was sent
        connector._send_message.assert_called()
        call_args = connector._send_message.call_args[0][0]
        assert call_args['event'] == 'subscribe'
        assert 'XBT/USD' in call_args['pair']
        assert call_args['subscription']['name'] == 'book'
        
        # Test trade subscription
        await connector.subscribe_trades('ETHUSDT')
        
        assert 'ETHUSDT' in connector.subscriptions
        assert 'trades' in connector.subscriptions['ETHUSDT']
        
        # Test unsubscription
        await connector.unsubscribe_orderbook('BTCUSDT')
        
        # Verify unsubscription message
        call_args = connector._send_message.call_args[0][0]
        assert call_args['event'] == 'unsubscribe'
    
    @pytest.mark.asyncio
    async def test_orderbook_snapshot_parsing(self, connector):
        """Test parsing order book snapshot data."""
        # Mock order book data from Kraken
        mock_data = {
            'bids': [
                ['50000.00', '1.5', 1609459200],
                ['49999.00', '2.0', 1609459201]
            ],
            'asks': [
                ['50001.00', '1.2', 1609459200],
                ['50002.00', '0.8', 1609459201]
            ]
        }
        
        # Parse the data
        orderbook = connector._parse_orderbook_snapshot(mock_data, 'BTCUSDT')
        
        # Verify parsing
        assert isinstance(orderbook, OrderBookSnapshot)
        assert orderbook.symbol == 'BTCUSDT'
        assert orderbook.exchange == 'kraken'
        
        # Verify bids
        assert len(orderbook.bids) == 2
        assert orderbook.bids[0].price == 50000.00
        assert orderbook.bids[0].size == 1.5
        assert orderbook.bids[1].price == 49999.00
        assert orderbook.bids[1].size == 2.0
        
        # Verify asks
        assert len(orderbook.asks) == 2
        assert orderbook.asks[0].price == 50001.00
        assert orderbook.asks[0].size == 1.2
        assert orderbook.asks[1].price == 50002.00
        assert orderbook.asks[1].size == 0.8
    
    @pytest.mark.asyncio
    async def test_orderbook_update_handling(self, connector):
        """Test handling order book update messages."""
        # Mock callback
        callback_called = False
        received_data = None
        
        def mock_callback(data):
            nonlocal callback_called, received_data
            callback_called = True
            received_data = data
        
        connector.add_data_callback(mock_callback)
        
        # Mock Kraken order book update message
        update_message = [
            123,  # channel ID
            {
                'b': [['50000.00', '1.5', '1609459200.123456']],
                'a': [['50001.00', '1.2', '1609459200.123456']]
            },
            'book-25',
            'XBT/USD'
        ]
        
        # Handle the message
        await connector._handle_orderbook_update(update_message)
        
        # Verify callback was called
        assert callback_called
        assert isinstance(received_data, OrderBookSnapshot)
        assert received_data.symbol == 'BTCUSDT'
        assert received_data.exchange == 'kraken'
        
        # Verify channel mapping was stored
        assert 123 in connector.channel_map
        assert connector.channel_map[123] == ('book-25', 'BTCUSDT')
        
        # Verify bids and asks
        assert len(received_data.bids) == 1
        assert received_data.bids[0].price == 50000.00
        assert received_data.bids[0].size == 1.5
        
        assert len(received_data.asks) == 1
        assert received_data.asks[0].price == 50001.00
        assert received_data.asks[0].size == 1.2
    
    @pytest.mark.asyncio
    async def test_trade_handling(self, connector):
        """Test handling trade messages."""
        # Mock callback
        callback_called = False
        received_trades = []
        
        def mock_callback(data):
            nonlocal callback_called
            callback_called = True
            received_trades.append(data)
        
        connector.add_data_callback(mock_callback)
        
        # Mock Kraken trade message (array of trades)
        trade_message = [
            456,  # channel ID
            [
                ['50000.50', '0.1', 1609459200.123456, 'b', 'm', ''],
                ['50001.00', '0.05', 1609459201.123456, 's', 'l', '']
            ],
            'trade',
            'XBT/USD'
        ]
        
        # Handle the message
        await connector._handle_trade_update(trade_message)
        
        # Verify callback was called
        assert callback_called
        assert len(received_trades) == 2
        
        # Verify first trade (buy)
        trade1 = received_trades[0]
        assert isinstance(trade1, TradeEvent)
        assert trade1.symbol == 'BTCUSDT'
        assert trade1.exchange == 'kraken'
        assert trade1.price == 50000.50
        assert trade1.size == 0.1
        assert trade1.side == 'buy'
        
        # Verify second trade (sell)
        trade2 = received_trades[1]
        assert trade2.price == 50001.00
        assert trade2.size == 0.05
        assert trade2.side == 'sell'
        
        # Verify channel mapping was stored
        assert 456 in connector.channel_map
        assert connector.channel_map[456] == ('trade', 'BTCUSDT')
    
    @pytest.mark.asyncio
    async def test_system_status_handling(self, connector):
        """Test handling system status messages."""
        # Mock system status message
        status_message = {
            'event': 'systemStatus',
            'status': 'online',
            'version': '1.0.0'
        }
        
        # Handle the message
        await connector._handle_system_status(status_message)
        
        # Verify status was updated
        assert connector.system_status == 'online'
    
    @pytest.mark.asyncio
    async def test_subscription_status_handling(self, connector):
        """Test handling subscription status messages."""
        # Mock subscription status message
        sub_message = {
            'event': 'subscriptionStatus',
            'status': 'subscribed',
            'channelName': 'book-25',
            'channelID': 123,
            'pair': 'XBT/USD',
            'subscription': {'name': 'book', 'depth': 25}
        }
        
        # Handle the message
        await connector._handle_subscription_status(sub_message)
        
        # Verify channel mapping was stored
        assert 123 in connector.channel_map
        assert connector.channel_map[123] == ('book-25', 'BTCUSDT')
    
    @pytest.mark.asyncio
    async def test_get_symbols(self, connector):
        """Test getting available symbols."""
        # Mock HTTP response
        mock_response_data = {
            'error': [],
            'result': {
                'XBTUSD': {
                    'wsname': 'XBT/USD'
                },
                'ETHUSD': {
                    'wsname': 'ETH/USD'
                },
                'XBTUSD.d': {  # Dark pool - should be filtered out
                    'wsname': 'XBT/USD.d'
                }
            }
        }
        
        with patch('aiohttp.ClientSession.get') as mock_get:
            mock_response = AsyncMock()
            mock_response.status = 200
            mock_response.json = AsyncMock(return_value=mock_response_data)
            mock_get.return_value.__aenter__.return_value = mock_response
            
            symbols = await connector.get_symbols()
            
            # Should only return non-dark pool symbols
            assert 'BTCUSDT' in symbols
            assert 'ETHUSDT' in symbols
            # Dark pool should be filtered out
            assert len([s for s in symbols if '.d' in s]) == 0
    
    @pytest.mark.asyncio
    async def test_get_orderbook_snapshot(self, connector):
        """Test getting order book snapshot from REST API."""
        # Mock HTTP response
        mock_orderbook = {
            'error': [],
            'result': {
                'XBTUSD': {
                    'bids': [['50000.00', '1.5', 1609459200]],
                    'asks': [['50001.00', '1.2', 1609459200]]
                }
            }
        }
        
        with patch('aiohttp.ClientSession.get') as mock_get:
            mock_response = AsyncMock()
            mock_response.status = 200
            mock_response.json = AsyncMock(return_value=mock_orderbook)
            mock_get.return_value.__aenter__.return_value = mock_response
            
            orderbook = await connector.get_orderbook_snapshot('BTCUSDT')
            
            assert isinstance(orderbook, OrderBookSnapshot)
            assert orderbook.symbol == 'BTCUSDT'
            assert orderbook.exchange == 'kraken'
            assert len(orderbook.bids) == 1
            assert len(orderbook.asks) == 1
    
    def test_authentication_token(self, connector):
        """Test authentication token generation."""
        # Set up credentials
        connector.api_key = 'test_key'
        connector.api_secret = 'dGVzdF9zZWNyZXQ='  # base64 encoded
        
        # Generate token
        token = connector._get_auth_token()
        
        # Should return a token (simplified implementation)
        assert isinstance(token, str)
        assert len(token) > 0
    
    def test_statistics(self, connector):
        """Test getting connector statistics."""
        # Add some test data
        connector.system_status = 'online'
        connector.channel_map[123] = ('book-25', 'BTCUSDT')
        
        stats = connector.get_kraken_stats()
        
        assert stats['exchange'] == 'kraken'
        assert stats['system_status'] == 'online'
        assert stats['channel_mappings'] == 1
        assert 'authenticated' in stats


async def test_kraken_integration():
    """Integration test for Kraken connector."""
    connector = KrakenConnector()
    
    try:
        # Test basic functionality
        assert connector.exchange_name == "kraken"
        
        # Test symbol normalization
        assert connector.normalize_symbol('BTCUSDT') == 'XBT/USD'
        assert connector._denormalize_symbol('XBT/USD') == 'BTCUSDT'
        
        # Test message type detection
        test_message = [123, {}, 'book-25', 'XBT/USD']
        assert connector._get_message_type(test_message) == 'book-25'
        
        # Test status message
        status_message = {'event': 'systemStatus', 'status': 'online'}
        assert connector._get_message_type(status_message) == 'systemStatus'
        
        print("✓ Kraken connector integration test passed")
        return True
        
    except Exception as e:
        print(f"✗ Kraken connector integration test failed: {e}")
        return False


if __name__ == "__main__":
    # Run integration test
    success = asyncio.run(test_kraken_integration())
    if not success:
        exit(1)