From b5c8c0b8da6a90996ac73ba052389524d6e6c3b5 Mon Sep 17 00:00:00 2001 From: Yuxuan Yan Date: Tue, 16 Jun 2026 17:42:20 +0800 Subject: [PATCH] Improve offline coverage for data boundaries --- tests/test_cli_workflow.py | 118 ++++++++++++++++++++++ tests/test_data_cli_and_resolve.py | 148 +++++++++++++++++++++++++++ tests/test_downloader_contracts.py | 156 +++++++++++++++++++++++++++++ tests/test_minute_downloader.py | 73 ++++++++++++++ tests/test_universe.py | 92 +++++++++++++++++ 5 files changed, 587 insertions(+) create mode 100644 tests/test_data_cli_and_resolve.py create mode 100644 tests/test_universe.py diff --git a/tests/test_cli_workflow.py b/tests/test_cli_workflow.py index e7d6b0a..a8f4d06 100644 --- a/tests/test_cli_workflow.py +++ b/tests/test_cli_workflow.py @@ -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 diff --git a/tests/test_data_cli_and_resolve.py b/tests/test_data_cli_and_resolve.py new file mode 100644 index 0000000..3335937 --- /dev/null +++ b/tests/test_data_cli_and_resolve.py @@ -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", + }] diff --git a/tests/test_downloader_contracts.py b/tests/test_downloader_contracts.py index 5797e81..1989ee8 100644 --- a/tests/test_downloader_contracts.py +++ b/tests/test_downloader_contracts.py @@ -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="用户未登录"), diff --git a/tests/test_minute_downloader.py b/tests/test_minute_downloader.py index 127e320..26b81c6 100644 --- a/tests/test_minute_downloader.py +++ b/tests/test_minute_downloader.py @@ -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", diff --git a/tests/test_universe.py b/tests/test_universe.py new file mode 100644 index 0000000..033c4f4 --- /dev/null +++ b/tests/test_universe.py @@ -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() == ["浦发银行", "华兴源创", "平安银行", "宁德时代"]