quant-trading/tests/test_live_pipeline.py

"""Contracts of the live pipeline: price panel refresh, session gating,
order sizing, and how results are reported."""

import unittest
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest import mock

import numpy as np
import pandas as pd

import bridge
import data_manager
import trader
from main import backtest


def _panel(rows: dict[str, list], start: str = "2026-09-14") -> pd.DataFrame:
    idx = pd.bdate_range(start, periods=len(next(iter(rows.values()))))
    return pd.DataFrame(rows, index=idx, dtype=float)


class PanelRefreshTests(unittest.TestCase):
    def test_refresh_overwrites_stale_window_and_keeps_missing_bars_raw(self):
        stored = _panel({"AAA": [10, 11, 11, 11, 12], "BBB": [5, 5, 5, 5, 5]})
        window_start = stored.index[2]
        fresh = pd.DataFrame(
            {"AAA": [11.5, np.nan, 12.2]}, index=stored.index[2:],
        )
        merged = data_manager._merge_refresh(stored, fresh, window_start)
        self.assertEqual(merged.loc[stored.index[2], "AAA"], 11.5)
        self.assertTrue(np.isnan(merged.loc[stored.index[3], "AAA"]))
        self.assertEqual(merged.loc[stored.index[4], "AAA"], 12.2)
        # Tickers not in the refresh keep their stored values.
        self.assertEqual(merged.loc[stored.index[3], "BBB"], 5)

    def test_refresh_rescales_history_after_split_adjustment(self):
        stored = _panel({"AAA": [100, 102, 104, 106, 108, 110, 112]})
        window_start = stored.index[2]
        # Yahoo now reports the whole window divided by 2 (2:1 split adjusted).
        fresh = stored.loc[window_start:, ["AAA"]] / 2
        merged = data_manager._merge_refresh(stored, fresh, window_start)
        returns = merged["AAA"].pct_change().dropna()
        self.assertTrue((returns > 0).all(), returns)
        self.assertAlmostEqual(merged["AAA"].iloc[0], 50.0)

    def test_newest_row_printed_by_few_tickers_is_dropped(self):
        panel = _panel({
            "A": [1, 1, 1, np.nan], "B": [1, 1, 1, np.nan],
            "C": [1, 1, 1, np.nan], "D": [1, 1, 1, 2],
        })
        self.assertNotIn(panel.index[-1], data_manager._drop_sparse_rows(panel).index)

    def test_rows_printed_by_a_minority_of_tickers_are_dropped(self):
        panel = _panel({
            "SPY": [1, 1, np.nan, 1, 1],
            "UPRO": [1, 1, np.nan, 1, 1],
            "GLD": [1, 1, np.nan, 1, 1],
            "7200.HK": [1, 1, 1, 1, 1],
        })
        cleaned = data_manager._drop_sparse_rows(panel)
        self.assertNotIn(panel.index[2], cleaned.index)
        self.assertEqual(len(cleaned), 4)


class LiveBookTests(unittest.TestCase):
    BOOK = {"label": "test", "start": {"us": "2026-09-16"}, "capital": {"us": 5_000.0},
            "fee_per_order": {"us": 1.0}, "strategies": {"us": ["momentum"]}}

    def _run(self, state: dict):
        close = _panel({"AAA": [10, 11, 12, 13, 14], "BBB": [20, 20, 21, 21, 22],
                        "SPY": [1, 1, 1, 1, 1]})
        args = SimpleNamespace(market="us", strategy="momentum", book="test", capital=None,
                               tx_cost=0.001, fixed_fee=2.0, integer_shares=True)
        equal = pd.DataFrame(0.5, index=close.index, columns=["AAA", "BBB"])
        with mock.patch.dict(trader.LIVE_BOOKS, {"test": self.BOOK}), \
                mock.patch.object(trader, "strategy_universe", lambda m, n: (["AAA", "BBB"], "SPY")), \
                mock.patch.object(trader, "pit_candidates", lambda m, n, s: []), \
                mock.patch.object(trader, "load_strategy_prices", lambda *a, **k: (close, close, close)), \
                mock.patch.object(trader, "session_signals", lambda m, n, t, days, c, o: equal.loc[days]):
            return trader._process_sessions(args, state, refresh=False)

    def test_a_new_account_starts_on_the_book_date_with_its_capital_and_flat_fee(self):
        state = self._run({})
        self.assertEqual(min(state["daily_equity"]), "2026-09-16")
        self.assertEqual(state["initial_capital"], 5_000.0)
        self.assertTrue(state["trade_log"])
        self.assertTrue(all(t["commission"] == 1.0 for t in state["trade_log"]))

    def test_no_account_is_opened_before_the_book_starts(self):
        with mock.patch.dict(self.BOOK, {"start": {"us": "2026-10-01"}}):
            self.assertIsNone(self._run({}))

    def test_book_strategies_are_live_strategies_of_their_market(self):
        for book, cfg in trader.LIVE_BOOKS.items():
            for market, names in cfg["strategies"].items():
                for name in names:
                    self.assertIn(name, trader.STRATEGY_REGISTRY, book)
                    self.assertNotIn(name, trader.LIVE_EXCLUDED_STRATEGIES, book)
                    self.assertTrue(trader.strategy_supports_market(name, market), (book, name))
                self.assertEqual(len(names), len(set(names)), book)
                pd.Timestamp(cfg["start"][market])
                self.assertGreater(cfg["capital"][market], 0)
                self.assertGreaterEqual(cfg["fee_per_order"][market], 0)


class SessionAndOrderTests(unittest.TestCase):
    def test_todays_bar_is_not_a_close_until_the_session_settles(self):
        raw = _panel({"SPY": [1, 1, 1]}, start="2026-09-22")
        today = raw.index[-1]
        intraday = datetime(2026, 9, 24, 17, 0, tzinfo=timezone.utc)   # 13:00 New York
        settled = datetime(2026, 9, 24, 20, 30, tzinfo=timezone.utc)   # 16:30 New York
        self.assertNotIn(today, trader.completed_sessions(raw, "us", "SPY", now=intraday))
        self.assertIn(today, trader.completed_sessions(raw, "us", "SPY", now=settled))

    def test_days_without_a_benchmark_print_are_not_sessions(self):
        raw = _panel({"SPY": [1, np.nan, 1], "7200.HK": [1, 1, 1]}, start="2026-09-02")
        later = datetime(2026, 10, 1, tzinfo=timezone.utc)
        sessions = trader.completed_sessions(raw, "etfs", "SPY", now=later)
        self.assertNotIn(raw.index[1], sessions)

    def test_a_share_orders_use_board_lots_and_exit_odd_lots(self):
        trades = trader.compute_trades(
            holdings={"600000.SS": 150}, cash=100_000.0,
            target_weights={"600519.SS": 0.5}, prices={"600519.SS": 170.0, "600000.SS": 10.0},
            min_trade_value=50, integer_shares=True, lot_size=100,
        )
        by_ticker = {t["ticker"]: t["shares_delta"] for t in trades}
        self.assertEqual(by_ticker["600519.SS"] % 100, 0)
        self.assertEqual(by_ticker["600000.SS"], -150)

    def test_tickers_without_a_fresh_price_are_not_traded(self):
        trades = trader.compute_trades(
            holdings={"UPRO": 10}, cash=5_000.0, target_weights={"TQQQ": 1.0},
            prices={"UPRO": 150.0, "TQQQ": 80.0}, frozen={"UPRO"},
        )
        self.assertEqual([t["ticker"] for t in trades], ["TQQQ"])


class ReportingTests(unittest.TestCase):
    def test_periodic_returns_include_the_first_day_of_each_period(self):
        idx = pd.to_datetime(["2026-01-30", "2026-02-02", "2026-02-27", "2026-03-02"])
        equity = pd.Series([100.0, 110.0, 110.0, 121.0], index=idx)
        monthly, yearly = bridge.compute_periodic_returns(equity)
        self.assertAlmostEqual(monthly["2026-02"], 0.10)
        self.assertAlmostEqual(monthly["2026-03"], 0.10)
        self.assertAlmostEqual(yearly["2026"], 0.21)

    def test_injected_cash_is_not_reported_as_performance(self):
        idx = pd.bdate_range("2024-01-02", periods=400)
        equity = pd.Series(10_000 * 1.0003 ** np.arange(400), index=idx)
        schedule = {"amount": 5_000, "frequency": "yearly_months", "months": [2, 8], "startYear": 1}
        plain = bridge._build_strategy_result("x", equity, [], None, 10_000)
        injected = bridge._build_strategy_result("x", equity, [], schedule, 10_000)
        self.assertAlmostEqual(injected["metrics"]["totalReturn"], plain["metrics"]["totalReturn"])
        flows = injected["capitalFlows"]
        self.assertEqual(flows["invested"], 10_000 + 5_000 * len(injected["injections"]))
        self.assertAlmostEqual(flows["profit"], flows["finalValue"] - flows["invested"], places=2)

    def test_ruined_portfolio_stays_at_zero(self):
        class TwiceLevered:
            def generate_signals(self, data):
                return pd.DataFrame(2.0, index=data.index, columns=data.columns)

        prices = pd.DataFrame({"X": [10.0, 10.0, 4.0, 5.0]}, index=pd.bdate_range("2026-01-01", periods=4))
        equity = backtest(TwiceLevered(), prices, initial_capital=100)
        self.assertTrue((equity.iloc[2:] == 0).all(), equity)


if __name__ == "__main__":
    unittest.main()