Yuxuan Yan
307 lines
8.7 KiB
Python
307 lines
8.7 KiB
Python
"""Offline downloader contract tests with mocked data providers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
import data.downloader as downloader
|
|
import pipeline.data.downloader as pipeline_downloader
|
|
from data.downloader import download_daily, download_daily_batch
|
|
from pipeline.common.schema import DATA_COLUMNS
|
|
from pipeline.data.downloader import download_universe
|
|
|
|
|
|
class _FakeResult:
|
|
def __init__(self, rows, error_code="0", error_msg=""):
|
|
self.rows = rows
|
|
self.error_code = error_code
|
|
self.error_msg = error_msg
|
|
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 _daily_batch_row(
|
|
date: str = "2024-01-02",
|
|
open_: str = "10",
|
|
high: str = "11",
|
|
low: str = "9",
|
|
close: str = "10.5",
|
|
preclose: str = "10",
|
|
volume: str = "1000",
|
|
amount: str = "10500",
|
|
) -> list[str]:
|
|
return [
|
|
date,
|
|
open_,
|
|
high,
|
|
low,
|
|
close,
|
|
preclose,
|
|
volume,
|
|
amount,
|
|
"1.23",
|
|
"5.0",
|
|
"1",
|
|
"0",
|
|
"8.0",
|
|
"1.1",
|
|
"2.2",
|
|
"3.3",
|
|
]
|
|
|
|
|
|
def test_download_daily_uses_baostock_before_akshare_in_auto(monkeypatch):
|
|
calls: list[str] = []
|
|
expected = pd.DataFrame({
|
|
"symbol": ["sh600000"],
|
|
"date": ["2024-01-02"],
|
|
"open": [10.0],
|
|
"high": [11.0],
|
|
"low": [9.0],
|
|
"close": [10.5],
|
|
"volume": [1000.0],
|
|
"amount": [10500.0],
|
|
})
|
|
|
|
def fake_baostock(symbol, start, end, adjust):
|
|
calls.append("baostock")
|
|
return expected
|
|
|
|
def fake_akshare(symbol, start, end, adjust):
|
|
calls.append("akshare")
|
|
raise AssertionError("akshare should not be called after baostock succeeds")
|
|
|
|
monkeypatch.setattr(downloader, "_download_baostock", fake_baostock)
|
|
monkeypatch.setattr(downloader, "_download_akshare", fake_akshare)
|
|
|
|
result = download_daily("sh600000", "2024-01-02", "2024-01-02", source="auto")
|
|
|
|
assert calls == ["baostock"]
|
|
assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]
|
|
assert result["close"].tolist() == [10.5]
|
|
|
|
|
|
def test_download_daily_falls_back_to_akshare_when_baostock_empty(monkeypatch):
|
|
calls: list[str] = []
|
|
fallback = pd.DataFrame({
|
|
"symbol": ["sz000001"],
|
|
"date": ["2024-01-02"],
|
|
"open": [20.0],
|
|
"high": [21.0],
|
|
"low": [19.0],
|
|
"close": [20.5],
|
|
"volume": [2000.0],
|
|
"amount": [41000.0],
|
|
})
|
|
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_download_baostock",
|
|
lambda symbol, start, end, adjust: calls.append("baostock") or None,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_download_akshare",
|
|
lambda symbol, start, end, adjust: calls.append("akshare") or fallback,
|
|
)
|
|
|
|
result = download_daily("sz000001", "2024-01-02", "2024-01-02", source="auto")
|
|
|
|
assert calls == ["baostock", "akshare"]
|
|
assert result["symbol"].tolist() == ["sz000001"]
|
|
assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]
|
|
|
|
|
|
def test_download_daily_batch_maps_rich_schema_and_vwap(monkeypatch):
|
|
query_calls: list[dict] = []
|
|
login_count = 0
|
|
logout_count = 0
|
|
|
|
def fake_login():
|
|
nonlocal login_count
|
|
login_count += 1
|
|
|
|
def fake_logout():
|
|
nonlocal logout_count
|
|
logout_count += 1
|
|
|
|
def fake_query(**kwargs):
|
|
query_calls.append(kwargs)
|
|
rows = [
|
|
_daily_batch_row(volume="1000", amount="10500"),
|
|
_daily_batch_row(date="2024-01-03", volume="0", amount="0"),
|
|
]
|
|
return _FakeResult(rows)
|
|
|
|
monkeypatch.setattr(downloader.bs, "login", fake_login)
|
|
monkeypatch.setattr(downloader.bs, "logout", fake_logout)
|
|
monkeypatch.setattr(downloader.bs, "query_history_k_data_plus", fake_query)
|
|
|
|
[(symbol, frame)] = list(
|
|
download_daily_batch(
|
|
["sh600000"],
|
|
"2024-01-02",
|
|
"2024-01-03",
|
|
adjust="hfq",
|
|
)
|
|
)
|
|
|
|
assert symbol == "sh600000"
|
|
assert query_calls[0]["code"] == "sh.600000"
|
|
assert query_calls[0]["adjustflag"] == "1"
|
|
assert login_count == 1
|
|
assert logout_count == 1
|
|
assert frame is not None
|
|
assert frame.columns.tolist() == [
|
|
"symbol", "date", "open", "high", "low", "close", "preclose",
|
|
"volume", "amount", "vwap", "turn", "pctChg", "tradestatus", "isST",
|
|
"peTTM", "pbMRQ", "psTTM", "pcfNcfTTM",
|
|
]
|
|
assert np.isclose(frame["vwap"].iloc[0], 10.5)
|
|
assert pd.isna(frame["vwap"].iloc[1])
|
|
assert pd.api.types.is_datetime64_any_dtype(frame["date"])
|
|
assert pd.api.types.is_numeric_dtype(frame["tradestatus"])
|
|
|
|
|
|
def test_download_daily_batch_relogs_and_retries_session_loss(monkeypatch):
|
|
responses = [
|
|
_FakeResult([], error_code="10002007", error_msg="用户未登录"),
|
|
_FakeResult([_daily_batch_row()]),
|
|
]
|
|
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),
|
|
)
|
|
|
|
[(symbol, frame)] = list(download_daily_batch(["sh600000"], "2024-01-02", "2024-01-02"))
|
|
|
|
assert symbol == "sh600000"
|
|
assert frame is not None
|
|
assert len(frame) == 1
|
|
assert login_count == 2
|
|
assert logout_count == 2
|
|
|
|
|
|
def test_download_daily_batch_uses_akshare_fallback_when_enabled(monkeypatch):
|
|
fallback = pd.DataFrame({
|
|
"symbol": ["sh600000"],
|
|
"date": ["2024-01-02"],
|
|
"open": [10.0],
|
|
"high": [11.0],
|
|
"low": [9.0],
|
|
"close": [10.5],
|
|
"volume": [1000.0],
|
|
"amount": [10500.0],
|
|
})
|
|
|
|
monkeypatch.setattr(downloader.bs, "login", lambda: None)
|
|
monkeypatch.setattr(downloader.bs, "logout", lambda: None)
|
|
monkeypatch.setattr(
|
|
downloader.bs,
|
|
"query_history_k_data_plus",
|
|
lambda **kwargs: _FakeResult([], error_code="1", error_msg="no data"),
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_download_akshare",
|
|
lambda symbol, start, end, adjust: fallback.copy(),
|
|
)
|
|
|
|
[(symbol, frame)] = list(
|
|
download_daily_batch(
|
|
["sh600000"],
|
|
"2024-01-02",
|
|
"2024-01-02",
|
|
akshare_fallback=True,
|
|
)
|
|
)
|
|
|
|
assert symbol == "sh600000"
|
|
assert frame is not None
|
|
assert frame["date"].tolist() == [pd.Timestamp("2024-01-02")]
|
|
assert frame["close"].tolist() == [10.5]
|
|
|
|
|
|
def test_download_universe_writes_daily_partitions_from_mock_batch(tmp_path, monkeypatch):
|
|
batch_frame = pd.DataFrame({
|
|
"symbol": ["sh600000", "sh600000"],
|
|
"date": pd.to_datetime(["2024-01-02", "2024-02-01"]),
|
|
"open": [10.0, 11.0],
|
|
"high": [11.0, 12.0],
|
|
"low": [9.0, 10.0],
|
|
"close": [10.5, 11.5],
|
|
"preclose": [10.0, 10.5],
|
|
"volume": [1000.0, 1200.0],
|
|
"amount": [10500.0, 13800.0],
|
|
"vwap": [10.5, 11.5],
|
|
"turn": [1.0, 1.1],
|
|
"pctChg": [5.0, 9.5],
|
|
"tradestatus": [1, 1],
|
|
"isST": [0, 0],
|
|
"peTTM": [8.0, 8.1],
|
|
"pbMRQ": [1.1, 1.2],
|
|
"psTTM": [2.1, 2.2],
|
|
"pcfNcfTTM": [3.1, 3.2],
|
|
})
|
|
|
|
monkeypatch.setattr(
|
|
pipeline_downloader,
|
|
"_resolve_universe",
|
|
lambda universe, max_symbols=0: pd.DataFrame({
|
|
"symbol_id": ["sh600000", "sz000001"],
|
|
"symbol_name": ["PF Bank", "Ping An Bank"],
|
|
}),
|
|
)
|
|
|
|
def fake_batch(symbols, start, end, adjust="qfq"):
|
|
assert symbols == ["sh600000", "sz000001"]
|
|
assert adjust == "qfq"
|
|
yield "sh600000", batch_frame
|
|
yield "sz000001", None
|
|
|
|
monkeypatch.setattr(pipeline_downloader, "download_daily_batch", fake_batch)
|
|
|
|
stats = download_universe(
|
|
universe="toy",
|
|
start_date="2024-01-02",
|
|
end_date="2024-02-01",
|
|
output_dir=str(tmp_path),
|
|
chunk_size=1,
|
|
)
|
|
|
|
dataset_path = tmp_path / "toy"
|
|
written = pd.read_parquet(dataset_path).sort_values(["date", "symbol_id"]).reset_index(drop=True)
|
|
assert stats == {
|
|
"dataset_path": str(dataset_path),
|
|
"n_symbols": 1,
|
|
"n_requested": 2,
|
|
"n_rows": 2,
|
|
"date_min": "2024-01-02",
|
|
"date_max": "2024-02-01",
|
|
}
|
|
assert (dataset_path / "month=2024-01").exists()
|
|
assert (dataset_path / "month=2024-02").exists()
|
|
assert written[DATA_COLUMNS].columns.tolist() == DATA_COLUMNS
|
|
assert written["symbol_name"].tolist() == ["PF Bank", "PF Bank"]
|