quant-trading/tests/test_execution_contract.py

import unittest
import json
import tempfile
from pathlib import Path
from unittest import mock

import numpy as np
import pandas as pd

import data_manager
import trader
from execution import simulate_book, rebalance, portfolio_value
from execution import compute_trades, market_permissions
from signals import session_targets
from strategies.momentum import MomentumStrategy
from strategies.momentum_quality import rolling_max_drawdown
from strategies.composite_alpha import CompositeAlphaStrategy
from strategies.robust_momentum import RobustMomentumStrategy


class ExecutionContracts(unittest.TestCase):
    def test_a_halt_cannot_be_sold_at_the_carried_mark(self):
        dates = pd.bdate_range("2026-01-05", periods=5)
        raw = pd.DataFrame({"A": [100, np.nan, np.nan, 10, 10]}, index=dates)
        targets = pd.DataFrame({"A": [1, 0, 0, 0, 0]}, index=dates)
        result = simulate_book(targets, raw, 10000, transaction_cost=0, min_trade_value=0)
        self.assertAlmostEqual(result["equity"].iloc[-1], 1000)
        self.assertEqual([t["date"] for t in result["trade_log"]],
                         [str(dates[0].date()), str(dates[3].date())])

    def test_fractional_and_integer_orders_are_funded_after_all_fees(self):
        prices = {"A": 73., "B": 29.}
        for integer in [False, True]:
            state = {"holdings": {}, "cash": 1000., "trade_log": []}
            rebalance(state, {"A": .5, "B": .5}, prices, set(prices),
                      fixed_fee=2, fee_base=1.88, fee_per_share=.009,
                      integer_shares=integer, min_trade_value=0)
            self.assertGreaterEqual(state["cash"], 0)
            equity = portfolio_value(state["holdings"], prices, state["cash"])
            self.assertAlmostEqual(equity, 1000 - sum(t["commission"] for t in state["trade_log"]), places=2)

    def test_batch_replay_matches_session_ledger_with_a_halt_and_drift(self):
        dates = pd.bdate_range("2026-01-05", periods=5)
        prices = pd.DataFrame({"A": [100, 110, np.nan, 70, 90], "B": [20, 21, 22, 22, 23]}, index=dates)
        targets = pd.DataFrame([[.5, .5], [.5, .5], [0, 1], [0, 1], [.5, .5]], index=dates, columns=prices.columns)
        options = dict(transaction_cost=.001, fixed_fee=2., integer_shares=True, resize_band=.1)
        batch = simulate_book(targets, prices, 10000, **options)
        state = trader.init_state("us", "test", 10000)
        fees = {"tx_cost": .001, "fixed_fee": 2.}
        for day in dates:
            trader._run_session(state, "us", day, "B", prices.ffill(), prices, fees, True, targets.loc[day])
            self.assertAlmostEqual(state["daily_equity"][str(day.date())], batch["equity"].loc[day], places=2)
        self.assertEqual(state["trade_log"], batch["trade_log"])

    def test_a_delisted_stock_is_settled_at_its_last_print_in_both_ledgers(self):
        dates = pd.bdate_range("2026-01-05", periods=9)
        # A is acquired: last print on day 2, out of the index from day 4.
        # H halts on the same day but stays a member, so it stays frozen.
        prices = pd.DataFrame({"A": [100, 101, 102] + [np.nan] * 6,
                               "H": [50, 50, 51] + [np.nan] * 6,
                               "B": [20.] * 9}, index=dates)
        member = pd.DataFrame(True, index=dates, columns=prices.columns)
        member.loc[dates[4]:, "A"] = False
        targets = pd.DataFrame([[.4, .3, .3]] * 9, index=dates, columns=prices.columns)
        options = dict(transaction_cost=.001, fixed_fee=2., integer_shares=True, resize_band=.1)
        batch = simulate_book(targets, prices, 10000, market="us", member=member, **options)
        settled = [(t["ticker"], t["date"], t["price"]) for t in batch["trade_log"] if t.get("settled")]
        self.assertEqual(settled, [("A", str(dates[7].date()), 102.0)])
        self.assertIn("H", batch["holdings"])
        state = trader.init_state("us", "test", 10000)
        fees = {"tx_cost": .001, "fixed_fee": 2.}
        for day in dates:
            trader._run_session(state, "us", day, "B", prices.ffill(), prices, fees, True, targets.loc[day],
                                member_row=member.loc[day])
            self.assertAlmostEqual(state["daily_equity"][str(day.date())], batch["equity"].loc[day], places=2)
        self.assertEqual(state["trade_log"], batch["trade_log"])
        china = simulate_book(targets, prices, 10000, market="cn", member=member, **options)
        self.assertFalse(any(t.get("settled") for t in china["trade_log"]))

    def test_contributions_are_executed_and_fees_affect_unit_returns(self):
        dates = pd.bdate_range("2026-01-05", periods=3)
        prices = pd.DataFrame({"A": [100., 100., 100.]}, index=dates)
        weights = pd.DataFrame(1., index=dates, columns=["A"])
        base = simulate_book(weights, prices, 1000, transaction_cost=0, fixed_fee=2)
        injected = simulate_book(weights, prices, 1000, transaction_cost=0, fixed_fee=2,
                                 contributions=pd.Series([0, 1000, 0], index=dates))
        self.assertEqual(len(base["trade_log"]), 1)
        self.assertEqual(len(injected["trade_log"]), 2)
        self.assertAlmostEqual(injected["equity"].iloc[-1], 1996)
        self.assertLess(injected["unit_equity"].iloc[-1], base["unit_equity"].iloc[-1])

    def test_missing_open_never_sees_same_day_close(self):
        close = pd.DataFrame({"A": [100., 200.]})
        raw_open = pd.DataFrame({"A": [np.nan, np.nan]})
        self.assertTrue(data_manager._fill_open(close, raw_open).isna().all().all())

    def test_china_limit_permissions_tax_and_star_minimum(self):
        raw = pd.Series({"600000.SS": 11., "000001.SZ": 9., "688001.SS": 12.})
        buy, sell, tax = market_permissions("cn", pd.Timestamp("2026-09-30"), raw, raw * 0 + 10)
        self.assertNotIn("600000.SS", buy)
        self.assertNotIn("000001.SZ", sell)
        self.assertNotIn("688001.SS", buy)
        self.assertEqual(tax, .0005)
        self.assertEqual(market_permissions("cn", pd.Timestamp("2023-08-25"), raw, raw * 0 + 10)[2], .001)
        orders = compute_trades({}, 10000, {"688001.SS": .7}, {"688001.SS": 23.},
                                integer_shares=True, lot_size=100)
        self.assertEqual(orders[0]["shares_delta"], 304)


class SignalContracts(unittest.TestCase):
    def test_live_refresh_updates_the_membership_snapshot_before_selection(self):
        with tempfile.TemporaryDirectory() as tmp:
            cache = Path(tmp) / "sp500_history.json"
            cache.write_text(json.dumps({"date": "2000-01-01", "intervals": {"OLD": [["2020-01-01", None]]}}))
            history = {"OLD": [["2020-01-01", "2026-09-30"]], "NEW": [["2026-09-30", None]]}
            with mock.patch.object(trader.universe_history, "CACHE_DIR", tmp), \
                    mock.patch.object(trader.universe_history, "build_sp500_history", return_value=history), \
                    mock.patch.object(trader, "strategy_universe", return_value=(["NEW"], "SPY")), \
                    mock.patch.object(data_manager, "update"):
                trader.refresh_panels("us", ["momentum"], "test")
                member = trader.index_members("us", pd.DatetimeIndex(["2026-10-01"]))
                self.assertTrue(member.iloc[0]["NEW"])
                self.assertFalse(member.iloc[0]["OLD"])

    def test_close_signal_is_unaffected_by_its_execution_days_close(self):
        dates = pd.bdate_range("2024-01-01", periods=30)
        c = pd.DataFrame({"A": np.arange(100., 130.), "B": np.arange(130., 100., -1)}, index=dates)
        s = MomentumStrategy(lookback=5, skip=1, top_n=1)
        base = session_targets(s, c)
        shocked = c.copy(); shocked.iloc[20:] = shocked.iloc[20:] * [0.1, 10]
        actual = session_targets(s, shocked)
        pd.testing.assert_frame_equal(base.iloc[:21], actual.iloc[:21])

    def test_trader_and_backtest_use_identical_membership_histories(self):
        rng = np.random.default_rng(42)
        dates = pd.bdate_range("2024-01-01", periods=370)
        close = pd.DataFrame(100 * np.exp(rng.normal(.0005, .02, (370, 25)).cumsum(0)),
                             index=dates, columns=[f"S{i}" for i in range(25)])
        opens = close.shift(1).fillna(100)
        member = pd.DataFrame(True, index=dates, columns=close.columns)
        member.iloc[:320, 0] = False
        member.iloc[330:, 1] = False
        bars = {"close": close, "open": opens}
        strategy = RobustMomentumStrategy("overnight_recovery", top_n=5)
        expected = session_targets(strategy, close, opens, bars, member)
        with mock.patch.object(trader, "index_members", lambda m, d: member.reindex(d)):
            actual = trader.session_signals("us", "overnight_recovery_top10", list(close), dates[-20:],
                                            close, opens, bars=bars, strategy=strategy)
        pd.testing.assert_frame_equal(actual, expected.reindex(index=actual.index, columns=actual.columns))
        self.assertTrue((actual["S1"] == 0).all())

    def test_quality_prefers_a_smaller_drawdown_and_needs_only_one_window(self):
        p = pd.DataFrame({"small": [100, 110, 99, 108], "large": [100, 110, 55, 108]})
        dd = rolling_max_drawdown(p, 4)
        self.assertGreater(dd.iloc[-1]["small"], dd.iloc[-1]["large"])
        self.assertAlmostEqual(dd.iloc[-1]["small"], -.1)

    def test_missing_regime_benchmark_is_an_error(self):
        with self.assertRaisesRegex(ValueError, "benchmark"):
            CompositeAlphaStrategy().generate_signals(pd.DataFrame({"A": [1, 2]}))


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