"""Tests for raw Baostock minute bar download plumbing.""" from pathlib import Path import numpy as np import pandas as pd import pytest import data.downloader as low_level_downloader import pipeline.data.downloader as pipeline_downloader from data.downloader import download_minute_batch from pipeline.common.schema import MINUTE_BAR_COLUMNS from pipeline.data.downloader import download_minute_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 test_download_minute_batch_maps_and_parses_baostock_rows(monkeypatch): rows = [ [ "2024-01-02", "20240102093500000", "sh.600000", "10", "11", "9", "10.5", "1000", "10500", "3", ], [ "2024-01-02", "20240102094000000", "sh.600000", "10.5", "12", "10", "11", "2000", "22000", "3", ], ] calls = [] def fake_query(**kwargs): calls.append(kwargs) return _FakeResult(rows) 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", fake_query, ) [(symbol, df)] = list( download_minute_batch( ["sh600000"], "2024-01-02", "2024-01-02", frequency=5, ) ) assert symbol == "sh600000" assert calls[0]["code"] == "sh.600000" assert calls[0]["frequency"] == "5" assert calls[0]["adjustflag"] == "3" assert df is not None assert df["datetime"].iloc[0] == pd.Timestamp("2024-01-02 09:35:00") assert df["time"].tolist() == ["09:35:00", "09:40:00"] assert (df["frequency"] == "5m").all() assert np.isclose(df["open"].iloc[0], 10.0) assert np.isclose(df["vwap"].iloc[0], 10.5) 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) monkeypatch.setattr( low_level_downloader.bs, "query_history_k_data_plus", lambda **kwargs: _FakeResult([]), ) assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [ ("sh600000", None) ] 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", "not-a-time", "sh.600000", "10", "11", "9", "10.5", "1000", "10500", "3", ]] 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: _FakeResult(bad_rows), ) assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [ ("sh600000", None) ] def test_download_minute_universe_writes_frequency_month_partitions(tmp_path, monkeypatch): minute = pd.DataFrame({ "symbol": ["sh600000", "sh600000"], "datetime": pd.to_datetime(["2024-01-02 09:35:00", "2024-01-02 09:40:00"]), "date": pd.to_datetime(["2024-01-02", "2024-01-02"]), "time": ["09:35:00", "09:40:00"], "frequency": ["5m", "5m"], "open": [10.0, 10.5], "high": [11.0, 12.0], "low": [9.0, 10.0], "close": [10.5, 11.0], "volume": [1000.0, 2000.0], "amount": [10500.0, 22000.0], "vwap": [10.5, 11.0], "adjustflag": ["3", "3"], }) monkeypatch.setattr( pipeline_downloader, "_resolve_universe", lambda universe, max_symbols=0: pd.DataFrame({ "symbol_id": ["sh600000"], "symbol_name": ["PF Bank"], }), ) def fake_batch(symbols, start, end, frequency=5): assert symbols == ["sh600000"] assert frequency == "5" yield "sh600000", minute monkeypatch.setattr(pipeline_downloader, "download_minute_batch", fake_batch) preserved = tmp_path / "toy" / "frequency=15m" / "month=2024-01" / "old.pq" preserved.parent.mkdir(parents=True) preserved_minute = minute.copy() preserved_minute["frequency"] = "15m" preserved_minute["symbol_id"] = "sh600000" preserved_minute["symbol_name"] = "PF Bank" preserved_minute[MINUTE_BAR_COLUMNS].to_parquet(preserved, index=False) stats = download_minute_universe( universe="toy", start_date="2024-01-02", end_date="2024-01-02", output_dir=str(tmp_path), chunk_size=1, frequency="5", ) dataset_path = Path(stats["dataset_path"]) assert (dataset_path / "frequency=5m" / "month=2024-01").is_dir() assert preserved.exists() out = pd.read_parquet(dataset_path / "frequency=5m") assert (set(MINUTE_BAR_COLUMNS) - {"frequency"}) <= set(out.columns) assert set(out["symbol_id"]) == {"sh600000"} assert set(out["symbol_name"]) == {"PF Bank"} assert stats["n_rows"] == 2 def test_download_minute_universe_raises_when_all_symbols_empty(tmp_path, monkeypatch): monkeypatch.setattr( pipeline_downloader, "_resolve_universe", lambda universe, max_symbols=0: pd.DataFrame({ "symbol_id": ["sh600000"], "symbol_name": ["PF Bank"], }), ) monkeypatch.setattr( pipeline_downloader, "download_minute_batch", lambda symbols, start, end, frequency=5: iter([("sh600000", None)]), ) with pytest.raises(RuntimeError, match="No minute data"): download_minute_universe( universe="toy", start_date="2024-01-02", end_date="2024-01-02", output_dir=str(tmp_path), )