yuxuanhui
367 lines
11 KiB
Python
367 lines
11 KiB
Python
"""Application tests for persisted whole-universe selection runs."""
|
|
|
|
import threading
|
|
import time
|
|
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 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):
|
|
return None
|
|
|
|
def get_latest_run(
|
|
self,
|
|
strategy: str,
|
|
target_trade_date: date | None = None,
|
|
*,
|
|
query: SelectionResultQuery | 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")
|
|
|
|
|
|
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
|