qlib/backend/tests/test_quant_engine.py

"""LocalEngine 回测测试：主路径、成本、涨跌停/不可买约束、无未来函数构造。"""

from __future__ import annotations

from datetime import date

import pandas as pd
from app.domain.entities.market import Stock
from app.domain.entities.research import (
    CostSpec,
    FactorSpec,
    ResearchSpec,
    SelectionSpec,
    UniverseSpec,
)
from app.quant.engine import LocalEngine
from app.quant.service import filter_stocks

from conftest_quant import synthetic_daily


def _spec(
    top_n: int = 1,
    start: date = date(2024, 3, 1),
    end: date = date(2024, 10, 31),
    rebalance: str = "monthly",
    costs: CostSpec | None = None,
) -> ResearchSpec:
    return ResearchSpec(
        type="backtest",
        universe=UniverseSpec(exclude_st=False, min_listing_days=0),
        factors=[FactorSpec(name="momentum_20")],
        selection=SelectionSpec(top_n=top_n),
        rebalance=rebalance,
        period=(start, end),
        costs=costs or CostSpec(),
    )


class TestBacktestMain:
    def test_uptrend_wins_and_profits(self) -> None:
        daily = synthetic_daily({"600000.SH": 0.004, "600001.SH": -0.002}, n=320)
        res = LocalEngine().run_backtest(daily, _spec(top_n=1))

        assert res.summary.total_return_pct > 0
        assert res.summary.final_equity > res.summary.initial_capital
        assert res.summary.total_trades >= 1
        assert res.summary.annual_return_pct > 0
        assert res.equity_curve[0].date == date(2024, 3, 1)
        assert res.equity_curve[-1].date == date(2024, 10, 31)
        assert res.monthly_returns
        assert res.trades
        # AGENT §24：未建模约束必须显式标注
        assert any("涨跌停" in item for item in res.unimplemented)
        assert res.config_snapshot["selection"]["top_n"] == 1

    def test_costs_reduce_returns(self) -> None:
        daily = synthetic_daily({"600000.SH": 0.004, "600001.SH": -0.002}, n=320)
        free = CostSpec(commission_rate=0.0, stamp_tax_rate=0.0, slippage_rate=0.0)
        with_cost = LocalEngine().run_backtest(daily, _spec(top_n=2, costs=CostSpec()))
        without = LocalEngine().run_backtest(daily, _spec(top_n=2, costs=free))
        # 有成本时收益不应高于无成本
        assert with_cost.summary.total_return_pct <= without.summary.total_return_pct + 1e-6

    def test_positions_and_drawdown_wellformed(self) -> None:
        daily = synthetic_daily({"600000.SH": 0.004, "600001.SH": -0.001}, n=260)
        res = LocalEngine().run_backtest(daily, _spec(top_n=2))
        assert res.positions
        weights = [p.weight for p in res.positions]
        assert all(0 < w <= 1 for w in weights)
        assert all(p.value <= 0 for p in res.drawdown)
        assert res.summary.max_drawdown_pct <= 0

    def test_required_columns_follow_factor_dependencies(self) -> None:
        """数据装配按引擎所需列裁剪：momentum 只取 close；量比因子需 volume。"""
        engine = LocalEngine()
        momentum = _spec()  # momentum_20
        assert engine.required_columns(momentum) == {"close"}
        volume = _spec().model_copy(update={"factors": [FactorSpec(name="volume_ratio_5_60")]})
        assert engine.required_columns(volume) == {"close", "volume"}
        unknown = _spec().model_copy(update={"factors": [FactorSpec(name="no_such")]})
        # 未知因子不参与列裁剪，交由执行期统一报错
        assert engine.required_columns(unknown) == {"close"}


def _limit_up_scenario_daily() -> pd.DataFrame:
    """Y 在 2024-07-01 相对前一交易日跳涨 10.5%（主板涨停不可追），且动量高于 X。"""
    dates = pd.bdate_range("2024-06-03", periods=46)
    rows = []
    for j, d in enumerate(dates):
        x = 100.0 * 1.001**j
        y = 100.0 * 1.003**j
        # 2024-07-01 是第 21 个工作日（6/28 周五 → 7/1 周一）
        if d.date() == date(2024, 7, 1):
            y = y / 1.003 * 1.105  # 相对前一日 +10.5%，形成涨停
        rows.append(
            {
                "symbol": "600001.SH",
                "trade_date": d.date(),
                "close": x,
                "open": x,
                "high": x * 1.01,
                "low": x * 0.99,
                "volume": 1e6,
                "amount": 1e8,
            }
        )
        rows.append(
            {
                "symbol": "600002.SH",
                "trade_date": d.date(),
                "close": y,
                "open": y,
                "high": y * 1.01,
                "low": y * 0.99,
                "volume": 1e6,
                "amount": 1e8,
            }
        )
    return pd.DataFrame(rows)


class TestLimitUpConstraint:
    def test_limit_up_symbol_not_bought(self) -> None:
        daily = _limit_up_scenario_daily()
        res = LocalEngine().run_backtest(
            daily, _spec(top_n=1, start=date(2024, 6, 3), end=date(2024, 8, 2))
        )
        # 7/1 调仓：Y 动量更高但涨停不可买 → 当日只能买入 X（8 月 Y 恢复可买，不断言之后）
        at_0701 = [p for p in res.positions if p.date == date(2024, 7, 1)]
        assert at_0701, "7/1 调仓后应记录持仓"
        assert {p.symbol for p in at_0701} == {"600001.SH"}


class TestServiceFilter:
    def _stock(self, symbol: str, name: str, list_date: date, delist: date | None = None) -> Stock:
        return Stock(symbol=symbol, name=name, list_date=list_date, delist_date=delist)

    def test_exclude_st_and_new_and_delisted(self) -> None:
        stocks = [
            self._stock("600001.SH", "*ST 某某", date(2000, 1, 1)),
            self._stock("600002.SH", "正常公司", date(2024, 6, 1)),  # 上市不足 250 天
            self._stock("600003.SH", "正常公司", date(2010, 1, 1), delist=date(2023, 6, 1)),
            self._stock("600004.SH", "正常公司", date(2010, 1, 1)),
        ]
        kept = filter_stocks(
            stocks, UniverseSpec(exclude_st=True, min_listing_days=250), as_of=date(2024, 8, 1)
        )
        assert [s.symbol for s in kept] == ["600004.SH"]