zhixing-system/zhixing-server/tests/unit/selection/test_run.py

"""Application tests for persisted whole-universe selection runs."""

from datetime import date
from decimal import Decimal
from typing import Literal

from zhixing_server.modules.selection.application.run import (
    PreparedSelectionRun,
    RunZhixingB1,
)
from zhixing_server.modules.selection.domain.models import (
    SelectionEvaluation,
    SelectionSignal,
    StockHistory,
)
from zhixing_server.modules.selection.domain.runs import (
    SelectionExecutionSource,
    SelectionResultQuery,
    SelectionRun,
    SelectionRunItem,
    SelectionRunStatus,
    SelectionStock,
)

TARGET = date(2026, 8, 8)


class FakeReader:
    def __init__(self, source: SelectionExecutionSource) -> None:
        self.source = source

    def load_execution_source(
        self,
        strategy: str,
        target_trade_date: date,
    ) -> SelectionExecutionSource:
        assert strategy == "zhixing_b1"
        assert target_trade_date == TARGET
        return self.source

    def load_history(self, ts_code: str, target_trade_date: date) -> StockHistory:
        raise AssertionError("the fake evaluator should be used")


class FakeStore:
    def __init__(self) -> None:
        self.items: list[SelectionRunItem] = []
        self.finished: tuple[str, SelectionRunStatus, dict[str, object]] | None = None

    def prepare_run(
        self,
        strategy: Literal["zhixing_b1"],
        target_trade_date: date,
        source: SelectionExecutionSource,
        *,
        rerun: bool,
    ) -> SelectionRun:
        assert strategy == "zhixing_b1"
        assert target_trade_date == TARGET
        assert rerun is False
        return SelectionRun(
            id="run-1",
            strategy="zhixing_b1",
            target_trade_date=TARGET,
            market_sync_batch_id=source.market_sync_batch_id,
            status="running",
            target_count=source.target_count,
            eligible_count=len(source.stocks),
            evaluated_count=0,
            selected_stock_count=0,
            signal_count=0,
            failed_count=0,
            coverage=source.coverage,
        )

    def record_item(self, run_id: str, item: SelectionRunItem) -> None:
        assert run_id == "run-1"
        self.items.append(item)

    def finish_run(
        self,
        run_id: str,
        status: SelectionRunStatus,
        *,
        evaluated_count: int,
        selected_stock_count: int,
        signal_count: int,
        failed_count: int,
        error_type: str | None = None,
        error_message: str | None = None,
    ) -> None:
        kwargs: dict[str, object] = {
            "evaluated_count": evaluated_count,
            "selected_stock_count": selected_stock_count,
            "signal_count": signal_count,
            "failed_count": failed_count,
        }
        if error_type is not None:
            kwargs["error_type"] = error_type
        if error_message is not None:
            kwargs["error_message"] = error_message
        self.finished = (run_id, status, kwargs)

    def get_run(self, run_id: str, *, query: SelectionResultQuery | None = None):
        return None

    def get_latest_run(
        self,
        strategy: str,
        target_trade_date: date | None = None,
        *,
        query: SelectionResultQuery | None = None,
    ):
        return None


class FakeEvaluator:
    def __init__(self, results: dict[str, SelectionEvaluation]) -> None:
        self.results = results

    def execute(self, ts_code: str, target_trade_date: date) -> SelectionEvaluation:
        return self.results[ts_code]


class RaisingEvaluator:
    def execute(self, ts_code: str, target_trade_date: date) -> SelectionEvaluation:
        if ts_code == "600000.SH":
            raise RuntimeError("temporary evaluator failure")
        return SelectionEvaluation(ts_code, target_trade_date, "no_signal")


def _source() -> SelectionExecutionSource:
    return SelectionExecutionSource(
        market_sync_batch_id="market-run-1",
        target_trade_date=TARGET,
        target_count=2,
        valid_count=2,
        coverage=Decimal("1"),
        stocks=(
            SelectionStock("000001.SZ", "平安银行"),
            SelectionStock("600000.SH", "浦发银行"),
        ),
    )


def _signal(ts_code: str, category: str) -> SelectionSignal:
    from zhixing_server.modules.selection.domain.models import ZhixingB1Category

    return SelectionSignal(
        ts_code=ts_code,
        name="平安银行",
        target_trade_date=TARGET,
        strategy="zhixing_b1",
        category=ZhixingB1Category(category),
        close=10.5,
        details={"j": 12.0},
    )


def test_prepare_captures_market_source_and_execute_persists_all_categories() -> None:
    source = _source()
    store = FakeStore()
    evaluator = FakeEvaluator(
        {
            "000001.SZ": SelectionEvaluation(
                "000001.SZ",
                TARGET,
                "selected",
                signals=(
                    _signal("000001.SZ", "zhixing_b1_original_b1"),
                    _signal("000001.SZ", "zhixing_b1_pullback_white"),
                ),
            ),
            "600000.SH": SelectionEvaluation(
                "600000.SH",
                TARGET,
                "no_signal",
                reason="no category matched",
            ),
        }
    )
    service = RunZhixingB1(FakeReader(source), store, evaluator)

    prepared = service.prepare("zhixing_b1", TARGET, rerun=False)
    assert isinstance(prepared, PreparedSelectionRun)
    service.execute(prepared)

    assert [item.status for item in store.items] == ["selected", "no_signal"]
    assert store.items[0].signal_count == 2
    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "success")
    assert store.finished[2] == {
        "evaluated_count": 2,
        "selected_stock_count": 1,
        "signal_count": 2,
        "failed_count": 0,
    }


def test_execute_marks_partial_success_when_one_stock_lacks_history() -> None:
    source = _source()
    store = FakeStore()
    evaluator = FakeEvaluator(
        {
            "000001.SZ": SelectionEvaluation("000001.SZ", TARGET, "no_signal"),
            "600000.SH": SelectionEvaluation(
                "600000.SH",
                TARGET,
                "insufficient_history",
                reason="warm-up data is incomplete",
            ),
        }
    )
    service = RunZhixingB1(FakeReader(source), store, evaluator)

    service.execute(service.prepare("zhixing_b1", TARGET, rerun=False))

    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "partial_success")
    assert store.finished[2]["failed_count"] == 1


def test_execute_isolates_unexpected_single_stock_failure() -> None:
    source = _source()
    store = FakeStore()
    service = RunZhixingB1(FakeReader(source), store, RaisingEvaluator())

    service.execute(service.prepare("zhixing_b1", TARGET, rerun=False))

    assert [item.status for item in store.items] == ["no_signal", "data_error"]
    assert store.items[1].reason == "temporary evaluator failure"
    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "partial_success")
    assert store.finished[2]["evaluated_count"] == 2
    assert store.finished[2]["failed_count"] == 1