tests/test_chart_colors.py

"""Tests for chart color consistency across all charts."""

import pytest

from src.folio.chart_data import (
    ChartColors,
    transform_for_allocations_chart,
    transform_for_exposure_chart,
    transform_for_treemap,
)
from src.folio.data_model import (
    ExposureBreakdown,
    OptionPosition,
    PortfolioGroup,
    PortfolioSummary,
    StockPosition,
)


class TestChartColors:
    """Tests for chart color consistency across all charts."""

    @pytest.fixture
    def mock_portfolio_summary(self):
        """Create a mock portfolio summary for testing."""
        # Create exposure breakdowns
        long_exposure = ExposureBreakdown(
            stock_exposure=10000.0,
            stock_beta_adjusted=12000.0,
            option_delta_exposure=2000.0,
            option_beta_adjusted=2400.0,
            total_exposure=12000.0,
            total_beta_adjusted=14400.0,
            description="Long market exposure (Stocks + Options)",
            formula="Long Stocks + Long Options Delta Exp",
            components={
                "Long Stocks Exposure": 10000.0,
                "Long Options Delta Exp": 2000.0,
                "Long Stocks Value": 10000.0,
                "Long Options Value": 2000.0,
            },
        )

        short_exposure = ExposureBreakdown(
            stock_exposure=-5000.0,  # Negative value
            stock_beta_adjusted=-6000.0,
            option_delta_exposure=-1000.0,  # Negative value
            option_beta_adjusted=-1200.0,
            total_exposure=-6000.0,
            total_beta_adjusted=-7200.0,
            description="Short market exposure (Stocks + Options)",
            formula="Short Stocks + Short Options Delta Exp",
            components={
                "Short Stocks Exposure": -5000.0,  # Negative value
                "Short Options Delta Exp": -1000.0,  # Negative value
                "Short Stocks Value": -5000.0,  # Negative value
                "Short Options Value": -1000.0,  # Negative value
            },
        )

        options_exposure = ExposureBreakdown(
            stock_exposure=0.0,
            stock_beta_adjusted=0.0,
            option_delta_exposure=1000.0,
            option_beta_adjusted=1200.0,
            total_exposure=1000.0,
            total_beta_adjusted=1200.0,
            description="Net delta exposure from options",
            formula="Long Options Delta Exp + Short Options Delta Exp (where Short is negative)",
            components={
                "Long Options Delta Exp": 2000.0,
                "Short Options Delta Exp": -1000.0,  # Negative value
                "Net Options Delta Exp": 1000.0,
            },
        )

        # Create portfolio summary
        return PortfolioSummary(
            net_market_exposure=6000.0,
            portfolio_beta=1.2,
            long_exposure=long_exposure,
            short_exposure=short_exposure,
            options_exposure=options_exposure,
            short_percentage=50.0,
            cash_like_positions=[],
            cash_like_value=3000.0,
            cash_like_count=1,
            cash_percentage=20.0,
            stock_value=5000.0,
            option_value=1000.0,
            pending_activity_value=500.0,
            portfolio_estimate_value=15000.0,
        )

    @pytest.fixture
    def mock_portfolio_groups(self):
        """Create mock portfolio groups for testing."""
        # Create stock positions
        aapl_stock = StockPosition(
            ticker="AAPL",
            position_type="stock",
            quantity=100,
            market_exposure=5000.0,
            beta=1.2,
            beta_adjusted_exposure=6000.0,
        )

        msft_stock = StockPosition(
            ticker="MSFT",
            position_type="stock",
            quantity=50,
            market_exposure=3000.0,
            beta=1.1,
            beta_adjusted_exposure=3300.0,
        )

        # Create option positions
        aapl_option = OptionPosition(
            ticker="AAPL",
            position_type="option",
            quantity=10,
            market_exposure=1000.0,
            beta=1.2,
            beta_adjusted_exposure=1200.0,
            strike=150.0,
            expiry="2023-01-01",
            option_type="CALL",
            delta=0.7,
            delta_exposure=700.0,
            notional_value=10000.0,
            underlying_beta=1.2,
        )

        # Create portfolio groups
        aapl_group = PortfolioGroup(
            ticker="AAPL",
            stock_position=aapl_stock,
            option_positions=[aapl_option],
            net_exposure=5700.0,
            beta=1.2,
            beta_adjusted_exposure=6840.0,
            total_delta_exposure=700.0,
            options_delta_exposure=700.0,
        )

        msft_group = PortfolioGroup(
            ticker="MSFT",
            stock_position=msft_stock,
            option_positions=[],
            net_exposure=3000.0,
            beta=1.1,
            beta_adjusted_exposure=3300.0,
            total_delta_exposure=0.0,
            options_delta_exposure=0.0,
        )

        return [aapl_group, msft_group]

    def test_color_constants_defined(self):
        """Test that all required color constants are defined in ChartColors."""
        # Core colors that should be defined
        assert hasattr(ChartColors, "LONG")
        assert hasattr(ChartColors, "SHORT")
        assert hasattr(ChartColors, "OPTIONS")
        assert hasattr(ChartColors, "NET")
        assert hasattr(ChartColors, "CASH")
        assert hasattr(ChartColors, "PENDING")

        # Verify they are all hex color strings
        for color_name in ["LONG", "SHORT", "OPTIONS", "NET", "CASH", "PENDING"]:
            color = getattr(ChartColors, color_name)
            assert isinstance(color, str)
            assert color.startswith("#")
            assert len(color) == 7  # #RRGGBB format

    def test_exposure_chart_uses_standard_colors(self, mock_portfolio_summary):
        """Test that the exposure chart uses the standard colors from ChartColors."""
        # Get the chart data
        chart_data = transform_for_exposure_chart(mock_portfolio_summary)

        # Extract the colors used in the chart
        colors = chart_data["data"][0]["marker"]["color"]

        # Verify that the colors match the ChartColors constants
        assert colors[0] == ChartColors.LONG  # Long
        assert colors[1] == ChartColors.SHORT  # Short
        assert colors[2] == ChartColors.OPTIONS  # Options
        assert colors[3] == ChartColors.NET  # Net

    def test_treemap_chart_uses_standard_colors(self, mock_portfolio_groups):
        """Test that the treemap chart uses the standard colors from ChartColors."""
        # Get the chart data
        chart_data = transform_for_treemap(mock_portfolio_groups)

        # Extract colors from the chart data
        # The first color is white for the root node, the rest are for the tickers
        colors = chart_data["data"][0]["marker"]["colors"]

        # We should have at least one green color (for AAPL and MSFT which are both long)
        # The first color is white for the root node
        assert colors[0] == "#FFFFFF"  # Root node is white

        # At least one of the other colors should be the LONG color
        assert ChartColors.LONG in colors[1:], "Long color not found in treemap"

        # SHORT color might not be present if there are no short positions in the test data
        # but we can verify the code would use it by checking the source
        import inspect

        source = inspect.getsource(transform_for_treemap)
        assert "ChartColors.LONG" in source
        assert "ChartColors.SHORT" in source

    def test_allocations_chart_uses_standard_colors(self, mock_portfolio_summary):
        """Test that the allocations chart uses the standard colors from ChartColors."""
        # Get the chart data
        chart_data = transform_for_allocations_chart(mock_portfolio_summary)

        # Extract the colors used in the chart
        long_color = chart_data["data"][0]["marker"]["color"]
        short_color = chart_data["data"][1]["marker"]["color"]
        cash_color = chart_data["data"][2]["marker"]["color"]
        pending_color = chart_data["data"][3]["marker"]["color"]

        # Verify that the colors match the ChartColors constants
        assert long_color == ChartColors.LONG
        assert short_color == ChartColors.SHORT
        assert cash_color == ChartColors.CASH
        assert pending_color == ChartColors.PENDING