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

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

import threading
import time
from collections.abc import Sequence
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.pattern_scoring import (
    PATTERN_SCORE_THRESHOLD,
    PATTERN_SCORING_VERSION,
    ZHIXING_B1_PATTERN_CASES,
    PatternCase,
    PatternScore,
    PatternScoreBreakdown,
)
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 BatchReader(FakeReader):
    def __init__(self, source: SelectionExecutionSource) -> None:
        super().__init__(source)
        self.batch_calls: list[tuple[str, ...]] = []

    def load_histories(
        self,
        stocks: tuple[SelectionStock, ...],
        target_trade_date: date,
    ) -> tuple[StockHistory, ...]:
        self.batch_calls.append(tuple(stock.ts_code for stock in stocks))
        return tuple(StockHistory(ts_code=stock.ts_code, name=stock.name) for stock in stocks)


class PartialBatchReader(BatchReader):
    def load_histories(
        self,
        stocks: tuple[SelectionStock, ...],
        target_trade_date: date,
    ) -> tuple[StockHistory, ...]:
        self.batch_calls.append(tuple(stock.ts_code for stock in stocks))
        return ()


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,
        sector_stock_codes: Sequence[str] | None = None,
    ):
        return None

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

    def get_run_identity(self, run_id: str):
        return None

    def get_latest_run_identity(self, strategy: str, target_trade_date: date | None = None):
        return None


class BatchStore(FakeStore):
    def __init__(self) -> None:
        super().__init__()
        self.batches: list[tuple[SelectionRunItem, ...]] = []

    def record_item(self, run_id: str, item: SelectionRunItem) -> None:
        raise AssertionError("the batch path should use record_items")

    def record_items(self, run_id: str, items: tuple[SelectionRunItem, ...]) -> None:
        assert run_id == "run-1"
        batch = tuple(items)
        self.batches.append(batch)
        self.items.extend(batch)


class FailingBatchStore(BatchStore):
    def record_items(self, run_id: str, items: tuple[SelectionRunItem, ...]) -> None:
        raise RuntimeError("batch write unavailable")


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")


class ConcurrentHistoryEvaluator:
    def __init__(self) -> None:
        self.active = 0
        self.peak = 0
        self._lock = threading.Lock()

    def execute(self, ts_code: str, target_trade_date: date) -> SelectionEvaluation:
        raise AssertionError("the batch evaluator path should be used")

    def execute_history(
        self,
        history: StockHistory,
        target_trade_date: date,
    ) -> SelectionEvaluation:
        with self._lock:
            self.active += 1
            self.peak = max(self.peak, self.active)
        time.sleep(0.02)
        with self._lock:
            self.active -= 1
        return SelectionEvaluation(history.ts_code, target_trade_date, "no_signal")


class FakePatternCaseLoader:
    def __init__(self, *, error: Exception | None = None) -> None:
        self.calls = 0
        self.error = error

    def load(self) -> tuple[PatternCase, ...]:
        self.calls += 1
        if self.error is not None:
            raise self.error
        return ()


class FakePatternScorer:
    def __init__(self, *, error: Exception | None = None) -> None:
        self.calls: list[str] = []
        self.error = error

    def score(self, history: StockHistory, cases: Sequence[PatternCase]) -> PatternScore:
        self.calls.append(history.ts_code)
        if self.error is not None:
            raise self.error
        return PatternScore(
            status="matched",
            value=88.0,
            threshold=PATTERN_SCORE_THRESHOLD,
            version=PATTERN_SCORING_VERSION,
            case=ZHIXING_B1_PATTERN_CASES[0],
            breakdown=PatternScoreBreakdown(80.0, 85.0, 90.0, 88.0),
        )


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


def test_execute_batches_history_reads_writes_and_limits_evaluation_workers() -> None:
    source = SelectionExecutionSource(
        market_sync_batch_id="market-run-1",
        target_trade_date=TARGET,
        target_count=8,
        valid_count=8,
        coverage=Decimal("1"),
        stocks=tuple(SelectionStock(f"{index:06d}.SZ", f"stock-{index}") for index in range(8)),
    )
    reader = BatchReader(source)
    store = BatchStore()
    evaluator = ConcurrentHistoryEvaluator()

    service = RunZhixingB1(
        reader,
        store,
        evaluator,
        max_workers=4,
        batch_size=4,
    )
    service.execute(service.prepare("zhixing_b1", TARGET, rerun=False))

    assert reader.batch_calls == [
        ("000000.SZ", "000001.SZ", "000002.SZ", "000003.SZ"),
        ("000004.SZ", "000005.SZ", "000006.SZ", "000007.SZ"),
    ]
    assert [len(batch) for batch in store.batches] == [4, 4]
    assert evaluator.peak <= 4
    assert evaluator.peak >= 2
    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "success")


def test_execute_maps_missing_batch_history_without_a_single_stock_read() -> None:
    source = _source()
    reader = PartialBatchReader(source)
    store = FakeStore()

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

    assert reader.batch_calls == [("000001.SZ", "600000.SH")]
    assert [item.status for item in store.items] == ["missing_target_bar", "missing_target_bar"]
    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "failed")


def test_execute_marks_batch_write_failure_as_failed() -> None:
    source = _source()
    reader = BatchReader(source)
    store = FailingBatchStore()
    evaluator = ConcurrentHistoryEvaluator()

    service = RunZhixingB1(reader, store, evaluator)
    service.execute(service.prepare("zhixing_b1", TARGET, rerun=False))

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


def test_execute_loads_cases_once_and_scores_only_selected_stocks() -> None:
    source = _source()
    reader = BatchReader(source)
    store = FakeStore()
    loader = FakePatternCaseLoader()
    scorer = FakePatternScorer()
    evaluator = FakeEvaluator(
        {
            "000001.SZ": SelectionEvaluation(
                "000001.SZ",
                TARGET,
                "selected",
                signals=(_signal("000001.SZ", "zhixing_b1_original_b1"),),
            ),
            "600000.SH": SelectionEvaluation("600000.SH", TARGET, "no_signal"),
        }
    )
    service = RunZhixingB1(
        reader,
        store,
        evaluator,
        pattern_case_loader=loader,
        pattern_scorer=scorer,
        pattern_scoring_enabled=True,
        batch_size=1,
    )

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

    assert loader.calls == 1
    assert scorer.calls == ["000001.SZ"]
    assert [item.pattern_score.status for item in store.items] == ["matched", "not_executed"]
    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "success")
    assert store.finished[2]["failed_count"] == 0


def test_execute_isolates_pattern_scoring_failure_from_selection_status() -> None:
    source = _source()
    store = FakeStore()
    loader = FakePatternCaseLoader()
    scorer = FakePatternScorer(error=RuntimeError("FastDTW unavailable"))
    evaluator = FakeEvaluator(
        {
            "000001.SZ": SelectionEvaluation(
                "000001.SZ",
                TARGET,
                "selected",
                signals=(_signal("000001.SZ", "zhixing_b1_original_b1"),),
            ),
            "600000.SH": SelectionEvaluation("600000.SH", TARGET, "no_signal"),
        }
    )
    service = RunZhixingB1(
        BatchReader(source),
        store,
        evaluator,
        pattern_case_loader=loader,
        pattern_scorer=scorer,
        pattern_scoring_enabled=True,
    )

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

    assert store.items[0].status == "selected"
    assert store.items[0].pattern_score == PatternScore.failed("FastDTW unavailable")
    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "success")
    assert store.finished[2]["failed_count"] == 0


def test_execute_skips_pattern_dependencies_when_feature_flag_is_disabled() -> None:
    source = _source()
    store = FakeStore()
    loader = FakePatternCaseLoader(error=AssertionError("loader must not run"))
    scorer = FakePatternScorer(error=AssertionError("scorer must not run"))
    evaluator = FakeEvaluator(
        {
            "000001.SZ": SelectionEvaluation(
                "000001.SZ",
                TARGET,
                "selected",
                signals=(_signal("000001.SZ", "zhixing_b1_original_b1"),),
            ),
            "600000.SH": SelectionEvaluation("600000.SH", TARGET, "no_signal"),
        }
    )
    service = RunZhixingB1(
        BatchReader(source),
        store,
        evaluator,
        pattern_case_loader=loader,
        pattern_scorer=scorer,
        pattern_scoring_enabled=False,
    )

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

    assert loader.calls == 0
    assert scorer.calls == []
    assert [item.pattern_score.status for item in store.items] == [
        "not_executed",
        "not_executed",
    ]
    assert store.items[0].signal_count == 1
    assert store.finished is not None
    assert store.finished[2]["failed_count"] == 0


def test_execute_marks_scores_failed_when_case_library_is_unavailable() -> None:
    source = _source()
    store = FakeStore()
    loader = FakePatternCaseLoader(error=RuntimeError("case_011 requires 25 qfq rows"))
    scorer = FakePatternScorer()
    evaluator = FakeEvaluator(
        {
            "000001.SZ": SelectionEvaluation(
                "000001.SZ",
                TARGET,
                "selected",
                signals=(_signal("000001.SZ", "zhixing_b1_original_b1"),),
            ),
            "600000.SH": SelectionEvaluation("600000.SH", TARGET, "no_signal"),
        }
    )
    service = RunZhixingB1(
        BatchReader(source),
        store,
        evaluator,
        pattern_case_loader=loader,
        pattern_scorer=scorer,
        pattern_scoring_enabled=True,
    )

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

    assert loader.calls == 1
    assert scorer.calls == []
    assert store.items[0].status == "selected"
    assert store.items[0].pattern_score.status == "failed"
    assert store.items[0].signals[0].category.value == "zhixing_b1_original_b1"
    assert store.finished is not None
    assert store.finished[0:2] == ("run-1", "success")
    assert store.finished[2]["failed_count"] == 0