Improve offline coverage for data boundaries

2026-06-16 17:42:20 +08:00
parent 31baa18ce5
commit b5c8c0b8da
5 changed files with 587 additions and 0 deletions
@@ -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",
+    }]