Add minute bar feature pipeline

tests/test_minute_downloader.py

@@ -0,0 +1,215 @@
+"""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_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_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),
+        )