Simple-Stock-Bot/tests/test_router.py

"""Test the Router symbol detection functionality."""

from unittest.mock import Mock, patch
from common.symbol_router import Router
from common.Symbol import Stock, Coin


class TestRouter:
    """Test Router symbol detection and parsing."""

    def setup_method(self):
        """Set up test fixtures."""
        # Mock the MarketData and cg_Crypto dependencies
        with patch("common.symbol_router.MarketData") as mock_md, patch("common.symbol_router.cg_Crypto") as mock_crypto:
            # Mock MarketData
            mock_md_instance = Mock()
            mock_md.return_value = mock_md_instance

            # Mock cg_Crypto
            mock_crypto_instance = Mock()
            mock_crypto.return_value = mock_crypto_instance

            # Create mock symbol list for crypto
            import pandas as pd

            mock_crypto_instance.symbol_list = pd.DataFrame(
                {
                    "symbol": ["btc", "eth", "ada"],
                    "id": ["bitcoin", "ethereum", "cardano"],
                    "name": ["Bitcoin", "Ethereum", "Cardano"],
                }
            )

            self.router = Router()
            self.mock_stock = mock_md_instance
            self.mock_crypto = mock_crypto_instance

    def test_stock_symbol_regex(self):
        """Test that stock symbol regex correctly identifies patterns."""
        import re

        # Test the regex pattern directly
        pattern = self.router.STOCK_REGEX

        # Should match
        assert re.search(pattern, "$AAPL")
        assert re.search(pattern, "What about $TSLA today?")
        assert re.search(pattern, "Look at $NVDA and $AMD")

        # Should not match crypto patterns
        assert not re.search(pattern, "$$BTC")

        # Should not match without $
        assert not re.search(pattern, "AAPL")

    def test_crypto_symbol_regex(self):
        """Test that crypto symbol regex correctly identifies patterns."""
        import re

        pattern = self.router.CRYPTO_REGEX

        # Should match
        assert re.search(pattern, "$$BTC")
        assert re.search(pattern, "How is $$ETH doing?")
        assert re.search(pattern, "$$BTC and $$ETH")

        # Should not match single $ patterns
        assert not re.search(pattern, "$BTC")
        assert not re.search(pattern, "$AAPL")

    def test_find_symbols_stock(self):
        """Test stock symbol detection."""
        # Mock successful stock lookup
        self.mock_stock.symbol_id.return_value = {"ticker": "AAPL", "title": "Apple Inc.", "mkt_cap_rank": 1}

        result = self.router.find_symbols("$AAPL")

        assert len(result) == 1
        assert isinstance(result[0], Stock)
        self.mock_stock.symbol_id.assert_called_with("AAPL")

    def test_find_symbols_crypto(self):
        """Test crypto symbol detection."""
        # Mock crypto dataframe lookup
        import pandas as pd

        mock_crypto_df = pd.DataFrame({"symbol": ["btc"], "id": ["bitcoin"], "name": ["Bitcoin"], "type_id": ["$$btc"]})

        # Set up the mock to return the DataFrame when filtered
        self.mock_crypto.symbol_list = pd.DataFrame(
            {
                "symbol": ["btc", "eth"],
                "id": ["bitcoin", "ethereum"],
                "name": ["Bitcoin", "Ethereum"],
                "type_id": ["$$btc", "$$eth"],
            }
        )

        # Mock the pandas filtering
        with patch.object(self.mock_crypto.symbol_list, "str") as mock_str:
            mock_str.fullmatch.return_value = mock_crypto_df
            mock_crypto_df.empty = False

            result = self.router.find_symbols("$$BTC")

            assert len(result) == 1
            assert isinstance(result[0], Coin)

    def test_find_symbols_no_matches(self):
        """Test behavior when no symbols are found."""
        result = self.router.find_symbols("Hello world")
        assert result == []

    def test_find_symbols_invalid_stock(self):
        """Test behavior when stock symbol is not found."""
        # Mock failed stock lookup
        self.mock_stock.symbol_id.return_value = None

        result = self.router.find_symbols("$INVALID")

        assert result == []
        self.mock_stock.symbol_id.assert_called_with("INVALID")

    def test_find_symbols_multiple(self):
        """Test finding multiple symbols."""
        # Mock successful lookups for both
        self.mock_stock.symbol_id.side_effect = [
            {"ticker": "AAPL", "title": "Apple Inc.", "mkt_cap_rank": 1},
            {"ticker": "TSLA", "title": "Tesla Inc.", "mkt_cap_rank": 5},
        ]

        result = self.router.find_symbols("$AAPL and $TSLA")

        assert len(result) == 2
        assert all(isinstance(symbol, Stock) for symbol in result)

    def test_trending_weight_tracking(self):
        """Test that trending counts are updated."""
        # Mock successful stock lookup
        self.mock_stock.symbol_id.return_value = {"ticker": "AAPL", "title": "Apple Inc.", "mkt_cap_rank": 1}

        # Clear trending count
        self.router.trending_count = {}

        result = self.router.find_symbols("$AAPL", trending_weight=2)

        assert len(result) == 1
        assert "$AAPL" in self.router.trending_count
        assert self.router.trending_count["$AAPL"] == 2