Improve offline coverage for data boundaries

This commit is contained in:
Yuxuan Yan
2026-06-16 17:42:20 +08:00
parent 31baa18ce5
commit b5c8c0b8da
5 changed files with 587 additions and 0 deletions
+118
View File
@@ -425,3 +425,121 @@ def test_cli_error_paths_are_clear_for_bad_user_inputs(tmp_path):
"--symbol", "sh999999",
])
assert "Symbol 'sh999999' not found" in alphaview_missing_symbol.output
def test_cli_list_and_legacy_feature_paths(tmp_path):
runner = CliRunner()
daily_bars = make_generated_daily_bars(n_sessions=3, include_missing=False)
minute_bars = make_generated_minute_bars(daily_bars)
daily_path = tmp_path / "daily_bars.pq"
minute_path = tmp_path / "minute_bars.pq"
daily_bars.to_parquet(daily_path, index=False)
minute_bars.to_parquet(minute_path, index=False)
derived_list = _invoke_ok(runner, ["derived", "list"])
feature_list = _invoke_ok(runner, ["feature", "list"])
assert "minute_daily_summary" in derived_list.output
assert "minute_daily_summary" in feature_list.output
feature_dir = tmp_path / "features"
feature_compute = _invoke_ok(runner, [
"feature",
"compute",
"--minute-path",
str(minute_path),
"--daily-path",
str(daily_path),
"--feature-type",
"minute_daily_summary",
"--feature-name",
"legacy_summary",
"--output-dir",
str(feature_dir),
])
feature_path = feature_dir / "legacy_summary.pq"
assert "Loaded minute bars:" in feature_compute.output
assert "Loaded daily data:" in feature_compute.output
assert "Saved feature:" in feature_compute.output
assert feature_path.exists()
no_input = _invoke_error(runner, [
"derived",
"compute",
"--derived-type",
"minute_daily_summary",
"--derived-name",
"missing_inputs",
"--output-dir",
str(tmp_path / "derived"),
])
assert "At least one of --daily-path or --minute-path is required" in no_input.output
missing_minute = _invoke_error(runner, [
"derived",
"compute",
"--daily-path",
str(daily_path),
"--derived-type",
"minute_daily_summary",
"--derived-name",
"daily_only",
"--output-dir",
str(tmp_path / "derived"),
])
assert "minute_daily_summary requires minute input" in missing_minute.output
malformed_feature_param = _invoke_error(runner, [
"feature",
"compute",
"--minute-path",
str(minute_path),
"--feature-type",
"minute_daily_summary",
"--feature-name",
"bad_param",
"--param",
"not-an-assignment",
"--output-dir",
str(tmp_path / "features_bad"),
])
assert "--param must be name=value" in malformed_feature_param.output
unknown_feature = _invoke_error(runner, [
"feature",
"compute",
"--minute-path",
str(minute_path),
"--feature-type",
"does_not_exist",
"--feature-name",
"bad_feature",
"--output-dir",
str(tmp_path / "features_unknown"),
])
assert "Unknown feature-type" in unknown_feature.output
def test_cli_pqcat_row_modes(tmp_path):
runner = CliRunner()
daily_bars = make_generated_daily_bars(n_sessions=3, include_missing=False)
daily_path = tmp_path / "daily_bars.pq"
daily_bars.to_parquet(daily_path, index=False)
head_result = _invoke_ok(runner, [
"pqcat",
str(daily_path),
"--head",
"2",
"--columns",
"symbol_id,close",
])
tail_result = _invoke_ok(runner, [
"pqcat",
str(daily_path),
"--tail",
"1",
])
assert "symbol_id" in head_result.output
assert "close" in head_result.output
assert "date" in tail_result.output
+148
View File
@@ -0,0 +1,148 @@
"""Offline coverage for data CLI and universe resolution glue."""
from __future__ import annotations
import pandas as pd
from click.testing import CliRunner
from cli import cli
import pipeline.data.cli as data_cli
import pipeline.data.downloader as pipeline_downloader
def test_resolve_universe_handles_named_file_all_and_symbol_list(tmp_path, monkeypatch):
hs300_raw = pd.DataFrame({
"updateDate": ["2024-01-12", "2024-01-12"],
"stockName": ["浦发银行", "平安银行"],
"stockCode": ["sh.600000", "sz.000001"],
})
zz500_raw = pd.DataFrame({
"name": ["东风汽车"],
"code": ["sh.600006"],
"date": ["2024-01-12"],
})
monkeypatch.setattr(pipeline_downloader, "get_hs300_stocks", lambda: hs300_raw)
monkeypatch.setattr(pipeline_downloader, "get_zz500_stocks", lambda: zz500_raw)
monkeypatch.setattr(
pipeline_downloader,
"get_all_stocks",
lambda: pd.DataFrame({
"code": ["sh600000", "sz000001", "sh600519"],
"name": ["浦发银行", "平安银行", "贵州茅台"],
}),
)
symbol_file = tmp_path / "symbols.txt"
symbol_file.write_text("sh600000\n\nsz000001\n")
hs300 = pipeline_downloader._resolve_universe("hs300")
zz500 = pipeline_downloader._resolve_universe("csi500")
all_capped = pipeline_downloader._resolve_universe("all", max_symbols=2)
from_file = pipeline_downloader._resolve_universe(str(symbol_file))
from_list = pipeline_downloader._resolve_universe("sh600000, sz000001")
assert hs300.to_dict("list") == {
"symbol_name": ["浦发银行", "平安银行"],
"symbol_id": ["sh600000", "sz000001"],
}
assert zz500.to_dict("list") == {
"symbol_name": ["东风汽车"],
"symbol_id": ["sh600006"],
}
assert all_capped["symbol_id"].tolist() == ["sh600000", "sz000001"]
assert from_file["symbol_id"].tolist() == ["sh600000", "sz000001"]
assert from_file["symbol_name"].tolist() == ["sh600000", "sz000001"]
assert from_list["symbol_id"].tolist() == ["sh600000", "sz000001"]
def test_data_cli_download_commands_print_summaries_without_network(monkeypatch, tmp_path):
runner = CliRunner()
daily_calls: list[dict] = []
minute_calls: list[dict] = []
def fake_daily(**kwargs):
daily_calls.append(kwargs)
return {
"dataset_path": str(tmp_path / "daily" / kwargs["universe"]),
"n_symbols": 2,
"n_requested": 3,
"n_rows": 18,
"date_min": "2024-01-02",
"date_max": "2024-01-12",
}
def fake_minute(**kwargs):
minute_calls.append(kwargs)
return {
"dataset_path": str(tmp_path / "minute" / kwargs["universe"]),
"frequency": "15m",
"n_symbols": 1,
"n_requested": 1,
"n_rows": 32,
"date_min": "2024-01-02",
"date_max": "2024-01-03",
}
monkeypatch.setattr(data_cli, "download_universe", fake_daily)
monkeypatch.setattr(data_cli, "download_minute_universe", fake_minute)
daily_result = runner.invoke(cli, [
"data",
"download",
"--universe",
"sh600000,sz000001",
"--start-date",
"2024-01-02",
"--end-date",
"2024-01-12",
"--output-dir",
str(tmp_path / "daily"),
"--symbols",
"3",
"--chunk-size",
"2",
"--adjust",
"none",
])
minute_result = runner.invoke(cli, [
"data",
"download-minute",
"--universe",
"toy",
"--start-date",
"2024-01-02",
"--end-date",
"2024-01-03",
"--output-dir",
str(tmp_path / "minute"),
"--symbols",
"1",
"--chunk-size",
"1",
"--frequency",
"15",
])
assert daily_result.exit_code == 0, daily_result.output
assert "Summary: 2/3 symbols, 18 bars" in daily_result.output
assert daily_calls == [{
"universe": "sh600000,sz000001",
"start_date": "2024-01-02",
"end_date": "2024-01-12",
"output_dir": str(tmp_path / "daily"),
"max_symbols": 3,
"chunk_size": 2,
"adjust": "none",
}]
assert minute_result.exit_code == 0, minute_result.output
assert "frequency=15m" in minute_result.output
assert minute_calls == [{
"universe": "toy",
"start_date": "2024-01-02",
"end_date": "2024-01-03",
"output_dir": str(tmp_path / "minute"),
"max_symbols": 1,
"chunk_size": 1,
"frequency": "15",
}]
+156
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import numpy as np
import pandas as pd
import pytest
import data.downloader as downloader
import pipeline.data.downloader as pipeline_downloader
@@ -119,6 +120,121 @@ def test_download_daily_falls_back_to_akshare_when_baostock_empty(monkeypatch):
assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]
def test_download_daily_raises_when_requested_source_has_no_data(monkeypatch):
monkeypatch.setattr(downloader, "_download_baostock", lambda *args: None)
with pytest.raises(RuntimeError, match="Failed to download data for sh600000"):
download_daily(
"sh600000",
"2024-01-02",
"2024-01-02",
source="baostock",
)
def test_akshare_daily_downloader_maps_columns_and_failures(monkeypatch):
calls: list[dict] = []
raw = pd.DataFrame({
"日期": ["2024-01-02"],
"开盘": [10.0],
"最高": [11.0],
"最低": [9.0],
"收盘": [10.5],
"成交量": [1000.0],
"成交额": [10500.0],
"换手率": [1.2],
})
def fake_hist(**kwargs):
calls.append(kwargs)
return raw.copy()
monkeypatch.setattr(downloader.ak, "stock_zh_a_hist", fake_hist)
result = downloader._download_akshare(
"sh600000",
"20240102",
"20240102",
adjust="",
)
assert calls == [{
"symbol": "600000",
"period": "daily",
"start_date": "20240102",
"end_date": "20240102",
"adjust": "",
}]
assert result is not None
assert result.columns.tolist() == [
"symbol", "date", "open", "high", "low", "close", "volume", "amount",
]
assert result["symbol"].tolist() == ["sh600000"]
monkeypatch.setattr(downloader.ak, "stock_zh_a_hist", lambda **kwargs: pd.DataFrame())
assert downloader._download_akshare("sh600000", "20240102", "20240102") is None
monkeypatch.setattr(
downloader.ak,
"stock_zh_a_hist",
lambda **kwargs: (_ for _ in ()).throw(RuntimeError("provider down")),
)
assert downloader._download_akshare("sh600000", "20240102", "20240102") is None
def test_baostock_daily_downloader_maps_errors_and_logout_failures(monkeypatch):
query_calls: list[dict] = []
row = ["2024-01-02", "10", "11", "9", "10.5", "1000", "10500"]
monkeypatch.setattr(downloader.bs, "login", lambda: None)
monkeypatch.setattr(
downloader.bs,
"logout",
lambda: (_ for _ in ()).throw(RuntimeError("logout failed")),
)
def fake_query(**kwargs):
query_calls.append(kwargs)
return _FakeResult([row])
monkeypatch.setattr(downloader.bs, "query_history_k_data_plus", fake_query)
result = downloader._download_baostock(
"sz000001",
"2024-01-02",
"2024-01-02",
adjust="none",
)
assert query_calls[0]["code"] == "sz.000001"
assert query_calls[0]["adjustflag"] == "3"
assert result is not None
assert result["symbol"].tolist() == ["sz000001"]
assert pd.api.types.is_numeric_dtype(result["close"])
monkeypatch.setattr(downloader.bs, "logout", lambda: None)
monkeypatch.setattr(
downloader.bs,
"query_history_k_data_plus",
lambda **kwargs: _FakeResult([], error_code="1", error_msg="bad symbol"),
)
assert downloader._download_baostock("sz000001", "2024-01-02", "2024-01-02") is None
monkeypatch.setattr(
downloader.bs,
"query_history_k_data_plus",
lambda **kwargs: _FakeResult([]),
)
assert downloader._download_baostock("sz000001", "2024-01-02", "2024-01-02") is None
monkeypatch.setattr(
downloader.bs,
"query_history_k_data_plus",
lambda **kwargs: (_ for _ in ()).throw(RuntimeError("query failed")),
)
assert downloader._download_baostock("sz000001", "2024-01-02", "2024-01-02") is None
def test_download_daily_batch_maps_rich_schema_and_vwap(monkeypatch):
query_calls: list[dict] = []
login_count = 0
@@ -170,6 +286,46 @@ def test_download_daily_batch_maps_rich_schema_and_vwap(monkeypatch):
assert pd.api.types.is_numeric_dtype(frame["tradestatus"])
def test_download_daily_batch_periodic_relogin_and_none_result(monkeypatch):
responses = [
_FakeResult([], error_code="1", error_msg="bad symbol"),
_FakeResult([_daily_batch_row(date="2024-01-03")]),
]
login_count = 0
logout_count = 0
def fake_login():
nonlocal login_count
login_count += 1
def fake_logout():
nonlocal logout_count
logout_count += 1
monkeypatch.setattr(downloader.bs, "login", fake_login)
monkeypatch.setattr(downloader.bs, "logout", fake_logout)
monkeypatch.setattr(
downloader.bs,
"query_history_k_data_plus",
lambda **kwargs: responses.pop(0),
)
results = list(
download_daily_batch(
["sh600000", "sz000001"],
"2024-01-02",
"2024-01-03",
relogin_every=1,
)
)
assert results[0] == ("sh600000", None)
assert results[1][0] == "sz000001"
assert results[1][1] is not None
assert login_count == 2
assert logout_count == 2
def test_download_daily_batch_relogs_and_retries_session_loss(monkeypatch):
responses = [
_FakeResult([], error_code="10002007", error_msg="用户未登录"),
+73
View File
@@ -91,6 +91,24 @@ def test_download_minute_batch_maps_and_parses_baostock_rows(monkeypatch):
assert pd.api.types.is_numeric_dtype(df["volume"])
def test_minute_frequency_and_timestamp_parsing_edge_cases():
frequency, label = low_level_downloader._normalize_minute_frequency("15m")
assert (frequency, label) == ("15", "15m")
with pytest.raises(ValueError, match="Unsupported minute frequency"):
low_level_downloader._normalize_minute_frequency("1m")
parsed = low_level_downloader._parse_minute_datetime(
pd.Series(["2024-01-02", "2024-01-02"]),
pd.Series(["0935", "09:40:00"]),
)
assert parsed.tolist() == [
pd.Timestamp("2024-01-02 09:35:00"),
pd.Timestamp("2024-01-02 09:40:00"),
]
def test_download_minute_batch_empty_result_yields_none(monkeypatch):
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
@@ -105,6 +123,61 @@ def test_download_minute_batch_empty_result_yields_none(monkeypatch):
]
def test_download_minute_batch_non_login_error_and_periodic_relogin(monkeypatch):
responses = [
_FakeResult([], error_code="1", error_msg="bad symbol"),
_FakeResult([]),
]
login_count = 0
logout_count = 0
def fake_login():
nonlocal login_count
login_count += 1
def fake_logout():
nonlocal logout_count
logout_count += 1
monkeypatch.setattr(low_level_downloader.bs, "login", fake_login)
monkeypatch.setattr(low_level_downloader.bs, "logout", fake_logout)
monkeypatch.setattr(
low_level_downloader.bs,
"query_history_k_data_plus",
lambda **kwargs: responses.pop(0),
)
assert list(
download_minute_batch(
["sh600000", "sz000001"],
"2024-01-02",
"2024-01-02",
relogin_every=1,
)
) == [("sh600000", None), ("sz000001", None)]
assert login_count == 2
assert logout_count == 2
def test_download_minute_batch_second_session_loss_yields_none(monkeypatch):
responses = [
_FakeResult([], error_code="10002007", error_msg="用户未登录"),
_FakeResult([], error_code="10002007", error_msg="用户未登录"),
]
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
monkeypatch.setattr(
low_level_downloader.bs,
"query_history_k_data_plus",
lambda **kwargs: responses.pop(0),
)
assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
("sh600000", None)
]
def test_download_minute_batch_rejects_unparsed_timestamps(monkeypatch):
bad_rows = [[
"2024-01-02",
+92
View File
@@ -0,0 +1,92 @@
"""Offline tests for baostock-backed universe helpers."""
from __future__ import annotations
import pandas as pd
import data.universe as universe
class _FakeResult:
def __init__(self, rows, fields=None):
self.rows = rows
self.fields = fields or ["code", "name", "date"]
self._idx = -1
def next(self):
self._idx += 1
return self._idx < len(self.rows)
def get_row_data(self):
return self.rows[self._idx]
def test_index_constituent_helpers_normalize_dotted_codes(monkeypatch):
calls: list[str] = []
monkeypatch.setattr(universe.bs, "login", lambda: calls.append("login"))
monkeypatch.setattr(universe.bs, "logout", lambda: calls.append("logout"))
monkeypatch.setattr(
universe.bs,
"query_hs300_stocks",
lambda: _FakeResult([
["sh.600000", "浦发银行", "2024-01-12"],
["sz.000001", "平安银行", "2024-01-12"],
]),
)
monkeypatch.setattr(
universe.bs,
"query_zz500_stocks",
lambda: _FakeResult([
["sh.600006", "东风汽车", "2024-01-12"],
]),
)
hs300 = universe.get_hs300_stocks()
zz500 = universe.get_zz500_stocks()
assert calls == ["login", "logout", "login", "logout"]
assert hs300["code"].tolist() == ["sh600000", "sz000001"]
assert zz500["code"].tolist() == ["sh600006"]
assert hs300["name"].tolist() == ["浦发银行", "平安银行"]
def test_get_all_stocks_walks_back_and_filters_to_listed_a_shares(monkeypatch):
fields = ["code", "tradeStatus", "code_name"]
responses = [
_FakeResult([], fields=fields),
_FakeResult(
[
["sh.600000", "1", "浦发银行"],
["sh.688001", "1", "华兴源创"],
["sz.000001", "1", "平安银行"],
["sz.300750", "1", "宁德时代"],
["sz.399001", "1", "深证成指"],
["sz.200001", "1", "深物业B"],
["bj.430047", "1", "北交所样本"],
],
fields=fields,
),
]
query_days: list[str] = []
monkeypatch.setattr(universe.bs, "login", lambda: None)
monkeypatch.setattr(universe.bs, "logout", lambda: None)
def fake_query_all_stock(day):
query_days.append(day)
return responses.pop(0)
monkeypatch.setattr(universe.bs, "query_all_stock", fake_query_all_stock)
result = universe.get_all_stocks("2024-01-07")
assert query_days == ["2024-01-07", "2024-01-06"]
assert result.columns.tolist() == ["code", "name"]
assert result["code"].tolist() == [
"sh600000",
"sh688001",
"sz000001",
"sz300750",
]
assert result["name"].tolist() == ["浦发银行", "华兴源创", "平安银行", "宁德时代"]