trading-agent/tests/test_guardrails.py

"""Pure-function tests for the risk guardrails — no network, deterministic."""

from __future__ import annotations

from datetime import datetime, timezone

from agents.schemas import AgentOutput, Pick
from config.risk import RiskConfig
from execution.guardrails import plan_orders

EQUITY = 100_000.0
CAPS = {"nova": 0.10, "warren_lite": 0.08, "atlas": 0.09}


def _out(persona_id: str, picks: list[Pick]) -> AgentOutput:
    return AgentOutput(
        agent_name=persona_id,
        persona_id=persona_id,
        started_at=datetime.now(timezone.utc),
        picks=picks,
    )


def _pick(ticker, conviction=0.7, weight=0.08, direction="long") -> Pick:
    return Pick(
        ticker=ticker, conviction=conviction, thesis="t",
        requested_weight_pct=weight, direction=direction,
    )


def _plan(agent_outputs, prices, *, risk=None, held=None, buying_power=EQUITY):
    return plan_orders(
        agent_outputs,
        equity=EQUITY,
        buying_power=buying_power,
        prices=prices,
        held_symbols=held or set(),
        risk=risk or RiskConfig(),
        persona_caps=CAPS,
    )


def test_shorts_dropped_when_disabled():
    orders, notes = _plan(
        [_out("nova", [_pick("NVDA"), _pick("TSLA", direction="short")])],
        {"NVDA": 200.0, "TSLA": 250.0},
        risk=RiskConfig(allow_shorts=False),
    )
    tickers = {o.ticker for o in orders}
    assert "NVDA" in tickers and "TSLA" not in tickers
    assert any("shorts disabled" in n for n in notes)


def test_shorts_kept_when_enabled():
    orders, _ = _plan(
        [_out("nova", [_pick("NVDA"), _pick("TSLA", direction="short")])],
        {"NVDA": 200.0, "TSLA": 250.0},
        risk=RiskConfig(allow_shorts=True),
    )
    by_ticker = {o.ticker: o for o in orders}
    assert "TSLA" in by_ticker and by_ticker["TSLA"].direction == "short"
    assert by_ticker["NVDA"].direction == "long"


def test_dedup_highest_conviction_wins():
    orders, _ = _plan(
        [
            _out("nova", [_pick("MSFT", conviction=0.6)]),
            _out("atlas", [_pick("MSFT", conviction=0.9)]),
        ],
        {"MSFT": 400.0},
    )
    msft = [o for o in orders if o.ticker == "MSFT"]
    assert len(msft) == 1
    assert msft[0].persona_id == "atlas" and msft[0].conviction == 0.9


def test_per_agent_trade_cap():
    picks = [_pick(t, conviction=c) for t, c in [("A", 0.9), ("B", 0.8), ("C", 0.7), ("D", 0.6)]]
    orders, notes = _plan(
        [_out("nova", picks)],
        {"A": 10.0, "B": 10.0, "C": 10.0, "D": 10.0},
        risk=RiskConfig(max_trades_per_agent=2),
    )
    assert len(orders) == 2
    assert {o.ticker for o in orders} == {"A", "B"}  # top conviction kept


def test_min_price_filter():
    orders, notes = _plan([_out("nova", [_pick("PENNY")])], {"PENNY": 1.5})
    assert orders == []
    assert any("min" in n for n in notes)


def test_skip_already_held():
    orders, notes = _plan(
        [_out("nova", [_pick("NVDA"), _pick("AMD")])],
        {"NVDA": 200.0, "AMD": 150.0},
        held={"NVDA"},
    )
    assert {o.ticker for o in orders} == {"AMD"}
    assert any("already held" in n for n in notes)


def test_sizing_and_weight_clamp():
    # requested 50% clamps to nova cap 10% -> qty = floor(0.10 * 100000 / 200) = 50
    orders, _ = _plan([_out("nova", [_pick("NVDA", weight=0.50)])], {"NVDA": 200.0})
    assert len(orders) == 1
    o = orders[0]
    assert abs(o.weight_pct - 0.10) < 1e-9
    assert o.qty == 50


def test_day_cap_truncates_by_conviction():
    picks = [_pick(f"T{i}", conviction=0.5 + i / 100) for i in range(6)]
    orders, notes = _plan(
        [_out("nova", picks[:3]), _out("atlas", picks[3:])],
        {f"T{i}": 100.0 for i in range(6)},
        risk=RiskConfig(max_trades_per_day=3, max_trades_per_agent=5),
    )
    assert len(orders) == 3
    assert any("day cap" in n for n in notes)


def test_option_sizing_by_premium():
    opt = Pick(
        ticker="AAPL", conviction=0.8, thesis="bullish", requested_weight_pct=0.10,
        instrument="option", option_right="call",
        contract_symbol="AAPL260724C00295000", strike=295.0, expiry="2026-07-24",
    )
    orders, _ = _plan([_out("nova", [opt])], {"AAPL260724C00295000": 7.39})
    assert len(orders) == 1
    o = orders[0]
    assert o.instrument == "option" and o.multiplier == 100
    assert o.contract_symbol == "AAPL260724C00295000"
    # qty = floor(0.10 * 100000 / (7.39 * 100)) = floor(10000 / 739) = 13
    assert o.qty == 13


def test_buying_power_trim():
    # 2 orders needing 50k each but only 60k buying power -> keep the higher conviction one
    orders, notes = _plan(
        [_out("nova", [_pick("HI", conviction=0.9, weight=0.10),
                       _pick("LO", conviction=0.5, weight=0.10)])],
        {"HI": 200.0, "LO": 200.0},
        buying_power=12_000.0,  # only one 50-share*200=10k order fits
    )
    assert len(orders) == 1
    assert orders[0].ticker == "HI"
    assert any("buying power" in n for n in notes)