tests/test_rehearsal.py

"""Tests for the PortfolioRehearsalEngine."""

import pytest
from unittest.mock import AsyncMock, MagicMock
from decimal import Decimal

from backend.agents.rehearsal import PortfolioRehearsalEngine


class TestPortfolioRehearsalEngine:
    """Tests for the PortfolioRehearsalEngine class."""

    @pytest.fixture
    def mock_mcp_router(self):
        """Create a mock MCP router."""
        router = MagicMock()
        router.call_yahoo_finance_mcp = AsyncMock(return_value={
            "close_prices": [100.0, 102.0, 101.0, 103.0, 105.0]
        })
        router.call_risk_analyzer_mcp = AsyncMock(return_value={
            "var_percentage": 0.05,
            "cvar_percentage": 0.07,
            "portfolio_volatility": 0.15,
            "sharpe_ratio": 1.2,
            "max_drawdown": 0.10
        })
        return router

    @pytest.fixture
    def engine(self, mock_mcp_router):
        """Create a PortfolioRehearsalEngine instance."""
        return PortfolioRehearsalEngine(mock_mcp_router)

    def test_engine_initialization(self, engine, mock_mcp_router):
        """Test engine initialisation."""
        assert engine.mcp_router == mock_mcp_router

    def test_apply_changes_sell(self, engine):
        """Test applying sell changes to portfolio."""
        portfolio = {
            "AAPL": {"weight": 50},
            "TSLA": {"weight": 30},
            "BND": {"weight": 20}
        }
        changes = [{"action": "sell", "ticker": "TSLA", "amount": 15}]

        result = engine._apply_changes(portfolio, changes)

        assert "AAPL" in result
        assert "TSLA" in result
        assert "BND" in result
        total = sum(h["weight"] for h in result.values())
        assert abs(total - 100) < 0.01

    def test_apply_changes_buy(self, engine):
        """Test applying buy changes to portfolio."""
        portfolio = {
            "AAPL": {"weight": 50},
            "BND": {"weight": 50}
        }
        changes = [{"action": "buy", "ticker": "TSLA", "amount": 20}]

        result = engine._apply_changes(portfolio, changes)

        assert "AAPL" in result
        assert "BND" in result
        assert "TSLA" in result
        total = sum(h["weight"] for h in result.values())
        assert abs(total - 100) < 0.01

    def test_apply_changes_sell_all(self, engine):
        """Test selling entire position removes ticker from portfolio."""
        portfolio = {
            "AAPL": {"weight": 50},
            "TSLA": {"weight": 50}
        }
        changes = [{"action": "sell", "ticker": "TSLA", "amount": 50}]

        result = engine._apply_changes(portfolio, changes)

        assert "AAPL" in result
        assert "TSLA" not in result
        assert result["AAPL"]["weight"] == 100

    def test_calculate_impact(self, engine):
        """Test impact calculation between portfolios."""
        current = {
            "expected_return": 10.0,
            "var_95": 5.0,
            "sharpe_ratio": 1.0,
            "num_holdings": 3
        }
        simulated = {
            "expected_return": 12.0,
            "var_95": 4.0,
            "sharpe_ratio": 1.2,
            "num_holdings": 4
        }

        impact = engine._calculate_impact(current, simulated)

        assert "expected_return" in impact
        assert impact["expected_return"]["current"] == 10.0
        assert impact["expected_return"]["simulated"] == 12.0
        assert impact["expected_return"]["delta"] == 2.0
        assert impact["expected_return"]["pct_change"] == 20.0

        assert impact["var_95"]["delta"] == -1.0

    def test_generate_recommendations_risk_reduction(self, engine):
        """Test recommendations for risk reduction."""
        impact = {
            "var_95": {"pct_change": -15},
            "expected_return": {"pct_change": 3},
            "sharpe_ratio": {"delta": 0.05},
            "diversification": {"delta": 0.02}
        }
        stress = {}

        recs = engine._generate_recommendations(impact, stress)

        assert any("Risk reduction" in rec for rec in recs)

    def test_generate_recommendations_risk_increase(self, engine):
        """Test recommendations for risk increase."""
        impact = {
            "var_95": {"pct_change": 15},
            "expected_return": {"pct_change": 3},
            "sharpe_ratio": {"delta": 0.05},
            "diversification": {"delta": 0.02}
        }
        stress = {}

        recs = engine._generate_recommendations(impact, stress)

        assert any("Risk increase" in rec for rec in recs)

    def test_assess_changes_recommended(self, engine):
        """Test assessment returns RECOMMENDED for positive changes."""
        impact = {
            "var_95": {"delta": -0.01},
            "expected_return": {"delta": 0.02},
            "sharpe_ratio": {"delta": 0.1},
            "diversification": {"delta": 0.05}
        }

        assessment = engine._assess_changes(impact)

        assert "RECOMMENDED" in assessment

    def test_assess_changes_caution(self, engine):
        """Test assessment returns CAUTION for negative changes."""
        impact = {
            "var_95": {"delta": 0.02},
            "expected_return": {"delta": -0.01},
            "sharpe_ratio": {"delta": -0.1},
            "diversification": {"delta": -0.05}
        }

        assessment = engine._assess_changes(impact)

        assert "CAUTION" in assessment

    def test_assess_changes_neutral(self, engine):
        """Test assessment returns NEUTRAL for mixed changes."""
        impact = {
            "var_95": {"delta": -0.01},
            "expected_return": {"delta": -0.01},
            "sharpe_ratio": {"delta": 0.1},
            "diversification": {"delta": -0.05}
        }

        assessment = engine._assess_changes(impact)

        assert "NEUTRAL" in assessment

    def test_estimate_beta(self, engine):
        """Test beta estimation for different ticker types."""
        assert engine._estimate_beta("TSLA", {}) == 1.5
        assert engine._estimate_beta("BND", {}) == 0.5
        assert engine._estimate_beta("AAPL", {}) == 1.0

    def test_calculate_expected_return(self, engine):
        """Test expected return calculation."""
        portfolio = {
            "AAPL": {"weight": 60},
            "BND": {"weight": 40}
        }
        historical_data = {
            "AAPL": {"close_prices": [100.0, 110.0]},
            "BND": {"close_prices": [50.0, 51.0]}
        }

        result = engine._calculate_expected_return(portfolio, historical_data)

        aapl_return = 10.0 * 0.6
        bnd_return = 2.0 * 0.4
        expected = aapl_return + bnd_return
        assert abs(result - expected) < 0.01

    @pytest.mark.asyncio
    async def test_rehearse_changes_structure(self, engine, mock_mcp_router):
        """Test that rehearse_changes returns expected structure."""
        current_portfolio = {
            "AAPL": {"weight": 50},
            "TSLA": {"weight": 30},
            "BND": {"weight": 20}
        }
        proposed_changes = [
            {"action": "sell", "ticker": "TSLA", "amount": 10},
            {"action": "buy", "ticker": "VTI", "amount": 10}
        ]

        result = await engine.rehearse_changes(
            current_portfolio,
            proposed_changes,
            portfolio_value=100000.0
        )

        assert "current" in result
        assert "simulated" in result
        assert "impact" in result
        assert "stress_comparison" in result
        assert "recommendations" in result
        assert "overall_assessment" in result

        assert "portfolio" in result["current"]
        assert "metrics" in result["current"]
        assert "portfolio" in result["simulated"]
        assert "metrics" in result["simulated"]

    @pytest.mark.asyncio
    async def test_rehearse_changes_empty_portfolio(self, engine, mock_mcp_router):
        """Test rehearsing changes with empty portfolio."""
        current_portfolio = {}
        proposed_changes = [
            {"action": "buy", "ticker": "VTI", "amount": 100}
        ]

        result = await engine.rehearse_changes(
            current_portfolio,
            proposed_changes,
            portfolio_value=100000.0
        )

        assert "current" in result
        assert "simulated" in result
        assert result["simulated"]["portfolio"]["VTI"]["weight"] == 100

    @pytest.mark.asyncio
    async def test_analyse_portfolio_empty(self, engine):
        """Test analysing empty portfolio returns zero metrics."""
        result = await engine._analyse_portfolio({}, {}, 100000.0)

        assert result["expected_return"] == 0.0
        assert result["var_95"] == 0.0
        assert result["num_holdings"] == 0

    @pytest.mark.asyncio
    async def test_compare_stress_tests(self, engine, mock_mcp_router):
        """Test stress test comparison."""
        current_portfolio = {"AAPL": {"weight": 100}}
        simulated_portfolio = {"BND": {"weight": 100}}
        historical_data = {}

        result = await engine._compare_stress_tests(
            current_portfolio,
            simulated_portfolio,
            historical_data,
            100000.0
        )

        assert "market_crash" in result
        assert "correction" in result
        assert "mild_decline" in result

        for scenario in result.values():
            assert "shock" in scenario
            assert "current_loss" in scenario
            assert "simulated_loss" in scenario
            assert "improvement" in scenario