path: root/services/order-executor/tests/test_risk_manager.py

"""Tests for RiskManager."""

from decimal import Decimal

from order_executor.risk_manager import RiskManager

from shared.models import OrderSide, Position, Signal


def make_signal(side: OrderSide, price: str, quantity: str, symbol: str = "AAPL") -> Signal:
    return Signal(
        strategy="test",
        symbol=symbol,
        side=side,
        price=Decimal(price),
        quantity=Decimal(quantity),
        reason="test signal",
    )


def make_risk_manager(
    max_position_size: str = "0.1",
    stop_loss_pct: str = "5.0",
    daily_loss_limit_pct: str = "10.0",
    trailing_stop_pct: str = "0",
    max_open_positions: int = 10,
    volatility_lookback: int = 20,
    volatility_scale: bool = False,
) -> RiskManager:
    return RiskManager(
        max_position_size=Decimal(max_position_size),
        stop_loss_pct=Decimal(stop_loss_pct),
        daily_loss_limit_pct=Decimal(daily_loss_limit_pct),
        trailing_stop_pct=Decimal(trailing_stop_pct),
        max_open_positions=max_open_positions,
        volatility_lookback=volatility_lookback,
        volatility_scale=volatility_scale,
    )


def make_position(symbol: str, quantity: str, avg_entry: str, current: str) -> Position:
    return Position(
        symbol=symbol,
        quantity=Decimal(quantity),
        avg_entry_price=Decimal(avg_entry),
        current_price=Decimal(current),
    )


def test_risk_check_passes_normal_order():
    """Small BUY order with enough balance should be allowed."""
    rm = make_risk_manager()
    signal = make_signal(side=OrderSide.BUY, price="100", quantity="0.5")
    # cost = 50, balance = 10000, position_value = 0 => (0+50)/10000 = 0.5% < 10%
    result = rm.check(signal, balance=Decimal("10000"), positions={}, daily_pnl=Decimal("0"))
    assert result.allowed is True
    assert result.reason == "OK"


def test_risk_check_rejects_exceeding_position_size():
    """5 BTC at $50,000 = $250,000 order cost on $10,000,000 balance exceeds 10% limit."""
    rm = make_risk_manager(max_position_size="0.1")
    signal = make_signal(side=OrderSide.BUY, price="50000", quantity="5")
    # cost = 250000, balance = 1000000 => 250000/1000000 = 25% > 10%
    # balance is sufficient (250000 < 1000000) but position size is exceeded
    result = rm.check(signal, balance=Decimal("1000000"), positions={}, daily_pnl=Decimal("0"))
    assert result.allowed is False
    assert result.reason == "Position size exceeded"


def test_risk_check_rejects_daily_loss_exceeded():
    """Daily PnL of -1100 on 10000 balance = -11%, exceeding -10% limit."""
    rm = make_risk_manager(daily_loss_limit_pct="10.0")
    signal = make_signal(side=OrderSide.BUY, price="100", quantity="0.1")
    result = rm.check(signal, balance=Decimal("10000"), positions={}, daily_pnl=Decimal("-1100"))
    assert result.allowed is False
    assert result.reason == "Daily loss limit exceeded"


def test_risk_check_rejects_insufficient_balance():
    """Order cost of 500 exceeds available balance of 100."""
    rm = make_risk_manager()
    signal = make_signal(side=OrderSide.BUY, price="100", quantity="5")
    # cost = 500, balance = 100
    result = rm.check(signal, balance=Decimal("100"), positions={}, daily_pnl=Decimal("0"))
    assert result.allowed is False
    assert result.reason == "Insufficient balance"


# --- Trailing stop tests ---


def test_trailing_stop_set_and_trigger():
    """Trailing stop should trigger when price drops below stop level."""
    rm = make_risk_manager(trailing_stop_pct="5")
    rm.set_trailing_stop("AAPL", Decimal("100"))

    signal = make_signal(side=OrderSide.BUY, price="94", quantity="0.01")
    result = rm.check(signal, balance=Decimal("10000"), positions={}, daily_pnl=Decimal("0"))
    assert result.allowed is False
    assert "Trailing stop triggered" in result.reason


def test_trailing_stop_updates_highest_price():
    """Trailing stop should track the highest price seen."""
    rm = make_risk_manager(trailing_stop_pct="5")
    rm.set_trailing_stop("AAPL", Decimal("100"))

    # Price rises to 120 => stop at 114
    rm.update_price("AAPL", Decimal("120"))

    # Price at 115 is above stop (114), should be allowed
    signal = make_signal(side=OrderSide.BUY, price="115", quantity="0.01")
    result = rm.check(signal, balance=Decimal("10000"), positions={}, daily_pnl=Decimal("0"))
    assert result.allowed is True

    # Price at 113 is below stop (114), should be rejected
    signal = make_signal(side=OrderSide.BUY, price="113", quantity="0.01")
    result = rm.check(signal, balance=Decimal("10000"), positions={}, daily_pnl=Decimal("0"))
    assert result.allowed is False
    assert "Trailing stop triggered" in result.reason


def test_trailing_stop_not_triggered_above_stop():
    """Trailing stop should not trigger when price is above stop level."""
    rm = make_risk_manager(trailing_stop_pct="5")
    rm.set_trailing_stop("AAPL", Decimal("100"))

    # Price at 96 is above stop (95), should be allowed
    signal = make_signal(side=OrderSide.BUY, price="96", quantity="0.01")
    result = rm.check(signal, balance=Decimal("10000"), positions={}, daily_pnl=Decimal("0"))
    assert result.allowed is True


# --- Max open positions test ---


def test_max_open_positions_check():
    """Should reject new BUY when max open positions is reached."""
    rm = make_risk_manager(max_open_positions=2)

    positions = {
        "AAPL": make_position("AAPL", "1", "100", "100"),
        "MSFT": make_position("MSFT", "10", "50", "50"),
    }

    signal = make_signal(side=OrderSide.BUY, price="10", quantity="1", symbol="TSLA")
    result = rm.check(signal, balance=Decimal("10000"), positions=positions, daily_pnl=Decimal("0"))
    assert result.allowed is False
    assert result.reason == "Max open positions reached"


# --- Volatility tests ---


def test_volatility_calculation():
    """Volatility should be calculated from price history."""
    rm = make_risk_manager(volatility_lookback=5)

    # No history yet
    assert rm.get_volatility("AAPL") is None

    # Feed prices
    prices = [100, 102, 98, 105, 101]
    for p in prices:
        rm.update_price("AAPL", Decimal(str(p)))

    vol = rm.get_volatility("AAPL")
    assert vol is not None
    assert vol > 0


def test_position_size_with_volatility_scaling():
    """High volatility should reduce position size."""
    rm = make_risk_manager(volatility_scale=True, volatility_lookback=5)

    # Feed volatile prices
    prices = [100, 120, 80, 130, 70]
    for p in prices:
        rm.update_price("AAPL", Decimal(str(p)))

    size = rm.calculate_position_size("AAPL", Decimal("10000"))
    base = Decimal("10000") * Decimal("0.1")

    # High volatility should reduce size below base
    assert size < base


def test_position_size_without_scaling():
    """Without scaling, position size should be base size regardless of volatility."""
    rm = make_risk_manager(volatility_scale=False, volatility_lookback=5)

    prices = [100, 120, 80, 130, 70]
    for p in prices:
        rm.update_price("AAPL", Decimal(str(p)))

    size = rm.calculate_position_size("AAPL", Decimal("10000"))
    base = Decimal("10000") * Decimal("0.1")

    assert size == base


# --- Portfolio exposure tests ---


def test_portfolio_exposure_check_passes():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        max_portfolio_exposure=0.8,
    )
    positions = {
        "AAPL": Position(
            symbol="AAPL",
            quantity=Decimal("0.01"),
            avg_entry_price=Decimal("50000"),
            current_price=Decimal("50000"),
        )
    }
    result = rm.check_portfolio_exposure(positions, Decimal("10000"))
    assert result.allowed  # 500/10000 = 5% < 80%


def test_portfolio_exposure_check_rejects():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        max_portfolio_exposure=0.3,
    )
    positions = {
        "AAPL": Position(
            symbol="AAPL",
            quantity=Decimal("1"),
            avg_entry_price=Decimal("50000"),
            current_price=Decimal("50000"),
        )
    }
    result = rm.check_portfolio_exposure(positions, Decimal("10000"))
    assert not result.allowed  # 50000/10000 = 500% > 30%


def test_correlation_calculation():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
    )
    # Feed identical price histories — correlation should be ~1.0
    for i in range(20):
        rm.update_price("A", Decimal(str(100 + i)))
        rm.update_price("B", Decimal(str(100 + i)))
    corr = rm.calculate_correlation("A", "B")
    assert corr is not None
    assert corr > 0.9


def test_var_calculation():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
    )
    for i in range(30):
        rm.update_price("AAPL", Decimal(str(100 + (i % 5) - 2)))
    positions = {
        "AAPL": Position(
            symbol="AAPL",
            quantity=Decimal("1"),
            avg_entry_price=Decimal("100"),
            current_price=Decimal("100"),
        )
    }
    var = rm.calculate_portfolio_var(positions, Decimal("10000"))
    assert var >= 0  # Non-negative


# --- Drawdown position scaling tests ---


def test_drawdown_position_scale_full():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        drawdown_reduction_threshold=0.1,
        drawdown_halt_threshold=0.2,
    )
    rm.update_balance(Decimal("10000"))
    scale = rm.get_position_scale(Decimal("10000"))
    assert scale == 1.0  # No drawdown


def test_drawdown_position_scale_reduced():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        drawdown_reduction_threshold=0.1,
        drawdown_halt_threshold=0.2,
    )
    rm.update_balance(Decimal("10000"))
    scale = rm.get_position_scale(Decimal("8500"))  # 15% drawdown (between 10% and 20%)
    assert 0.25 < scale < 1.0


def test_drawdown_halt():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        drawdown_reduction_threshold=0.1,
        drawdown_halt_threshold=0.2,
    )
    rm.update_balance(Decimal("10000"))
    scale = rm.get_position_scale(Decimal("7500"))  # 25% drawdown
    assert scale == 0.0


def test_consecutive_losses_pause():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        max_consecutive_losses=3,
        loss_pause_minutes=60,
    )
    rm.record_trade_result(False)
    rm.record_trade_result(False)
    assert not rm.is_paused()
    rm.record_trade_result(False)  # 3rd loss
    assert rm.is_paused()


def test_consecutive_losses_reset_on_win():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        max_consecutive_losses=3,
    )
    rm.record_trade_result(False)
    rm.record_trade_result(False)
    rm.record_trade_result(True)  # Win resets counter
    rm.record_trade_result(False)
    assert not rm.is_paused()  # Only 1 consecutive loss


def test_drawdown_check_rejects_in_check():
    rm = RiskManager(
        max_position_size=Decimal("0.5"),
        stop_loss_pct=Decimal("5"),
        daily_loss_limit_pct=Decimal("10"),
        drawdown_halt_threshold=0.15,
    )
    rm.update_balance(Decimal("10000"))
    signal = Signal(
        strategy="test",
        symbol="AAPL",
        side=OrderSide.BUY,
        price=Decimal("50000"),
        quantity=Decimal("0.01"),
        reason="test",
    )
    result = rm.check(signal, Decimal("8000"), {}, Decimal("0"))  # 20% dd > 15%
    assert not result.allowed
    assert "halted" in result.reason.lower()