trading-bot/tests/test_trade.py

from pathlib import Path
from typing import Any

import pandas as pd
import pytest

from trading_bot.models.prediction import PredictionResult
from trading_bot.models.trade import (
    rebalance_alpaca_portfolio,
    summarize_alpaca_portfolio,
)


def test_calculate_cash_to_spy_ratio() -> None:
    account = {"cash": "1000.00"}
    positions = [{"symbol": "SPY", "qty": "10", "market_value": "5000", "avg_entry_price": "500"}]
    orders = []

    summary = summarize_alpaca_portfolio(account, positions, orders)

    assert summary.cash == 1000.0
    assert summary.spy_quantity == 10.0
    assert summary.spy_market_value == 5000.0
    assert summary.current_cash_to_spy_ratio == 0.2
    assert summary.projected_cash_to_spy_ratio == 0.2
    assert summary.open_spy_order_count == 0


def test_projected_ratio_includes_open_spy_buy_order() -> None:
    account = {"cash": "1000.00"}
    positions = [{"symbol": "SPY", "qty": "10", "market_value": "5000", "avg_entry_price": "500"}]
    orders = [
        {
            "symbol": "SPY",
            "side": "buy",
            "qty": "2",
            "filled_qty": "0",
            "limit_price": "500",
            "status": "open",
            "type": "limit",
        }
    ]

    summary = summarize_alpaca_portfolio(account, positions, orders)

    assert summary.open_spy_order_count == 1
    assert summary.estimated_spy_price == 500.0
    assert summary.projected_cash_to_spy_ratio == 0.0


def test_projected_ratio_handles_no_spy_position_but_open_order() -> None:
    account = {"cash": "1500.00"}
    positions = []
    orders = [
        {
            "symbol": "SPY",
            "side": "buy",
            "qty": "3",
            "filled_qty": "0",
            "limit_price": "500",
            "status": "open",
            "type": "limit",
        }
    ]

    summary = summarize_alpaca_portfolio(account, positions, orders)

    assert summary.spy_quantity == 0.0
    assert summary.spy_market_value == 0.0
    assert summary.current_cash_to_spy_ratio is None
    assert summary.projected_cash_to_spy_ratio == 0.0
    assert summary.open_spy_order_count == 1


class FakeAlpacaClient:
    def __init__(self) -> None:
        self.cancelled = False
        self.order_request = None

    def get_account(self) -> dict[str, str]:
        return {"cash": "1000.00"}

    def get_all_positions(self) -> list[dict[str, str]]:
        return [
            {
                "symbol": "SPY",
                "qty": "10",
                "market_value": "5000",
                "avg_entry_price": "500",
            }
        ]

    def get_orders(self, filter=None) -> list[dict[str, Any]]:
        return []

    def cancel_orders(self) -> list[dict[str, Any]]:
        self.cancelled = True
        return []

    def submit_order(self, order_request: Any) -> dict[str, Any]:
        self.order_request = order_request
        return {
            "id": "fake-order",
            "status": "new",
            "symbol": getattr(order_request, "symbol", "SPY"),
            "side": getattr(order_request, "side", "buy"),
            "notional": getattr(order_request, "notional", 0.0),
        }


def test_rebalance_submits_buy_order_when_target_exposure_is_higher(monkeypatch) -> None:
    fake_client = FakeAlpacaClient()

    def fake_create_client(api_key=None, secret_key=None, paper=True):
        return fake_client

    def fake_predict_latest_probability(**kwargs):
        return PredictionResult(
            prediction_date=pd.Timestamp("2026-08-03"),
            probability=0.9,
            base_probability=0.5,
            feature_columns=["SPY_ret_5"],
        )

    monkeypatch.setattr(
        "trading_bot.models.trade.create_alpaca_trading_client",
        fake_create_client,
    )
    monkeypatch.setattr(
        "trading_bot.models.trade.predict_latest_probability",
        fake_predict_latest_probability,
    )

    summary, order = rebalance_alpaca_portfolio(
        api_key="test-key",
        secret_key="test-secret",
        paper=True,
        symbol="SPY",
        min_order_dollars=100.0,
        sensitivity=5.0,
    )

    assert fake_client.cancelled is True
    assert order is not None
    assert order["side"] == "buy"
    assert order["notional"] == pytest.approx(1000.0)
    assert summary.cash == 1000.0
    assert summary.spy_market_value == 5000.0


def test_rebalance_skips_small_orders_below_minimum(monkeypatch) -> None:
    fake_client = FakeAlpacaClient()

    def fake_create_client(api_key=None, secret_key=None, paper=True):
        return fake_client

    def fake_predict_latest_probability(**kwargs):
        return PredictionResult(
            prediction_date=pd.Timestamp("2026-08-03"),
            probability=0.569,
            base_probability=0.5,
            feature_columns=["SPY_ret_5"],
        )

    monkeypatch.setattr(
        "trading_bot.models.trade.create_alpaca_trading_client",
        fake_create_client,
    )
    monkeypatch.setattr(
        "trading_bot.models.trade.predict_latest_probability",
        fake_predict_latest_probability,
    )

    summary, order = rebalance_alpaca_portfolio(
        api_key="test-key",
        secret_key="test-secret",
        paper=True,
        symbol="SPY",
        min_order_dollars=100.0,
        sensitivity=5.0,
    )

    assert fake_client.cancelled is False
    assert order is None
    assert summary.cash == 1000.0
    assert summary.spy_market_value == 5000.0