chinese-equity-quant/tests/test_minute_downloader.py

"""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),
        )