Improve offline coverage for data boundaries
This commit is contained in:
Yuxuan Yan
@@ -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
|
||||
|
||||
@@ -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",
|
||||
}]
|
||||
@@ -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="用户未登录"),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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() == ["浦发银行", "华兴源创", "平安银行", "宁德时代"]
|
||||
Reference in New Issue
Block a user