Yuxuan Yan
149 lines
4.6 KiB
Python
149 lines
4.6 KiB
Python
"""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",
|
|
}]
|