chinese-equity-quant/tests/test_downloader_contracts.py

"""Offline downloader contract tests with mocked data providers."""

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
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_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_download_daily_akshare_source_skips_baostock(monkeypatch):
    calls: list[str] = []
    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,
        "_download_baostock",
        lambda *args: calls.append("baostock") or None,
    )
    monkeypatch.setattr(
        downloader,
        "_download_akshare",
        lambda *args: calls.append("akshare") or fallback,
    )

    result = download_daily(
        "sh600000",
        "2024-01-02",
        "2024-01-02",
        source="akshare",
    )

    assert calls == ["akshare"]
    assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]


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
    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_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_empty_rows_yields_none(monkeypatch):
    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([]),
    )

    assert list(download_daily_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
        ("sh600000", None)
    ]


def test_download_daily_batch_generic_exception_yields_none(monkeypatch):
    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: (_ for _ in ()).throw(RuntimeError("query failed")),
    )

    assert list(download_daily_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
        ("sh600000", None)
    ]


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_second_session_loss_and_logout_failure(monkeypatch):
    responses = [
        _FakeResult([], error_code="10002007", error_msg="用户未登录"),
        _FakeResult([], error_code="10002007", error_msg="用户未登录"),
    ]
    logout_count = 0

    def fake_logout():
        nonlocal logout_count
        logout_count += 1
        raise RuntimeError("logout failed")

    monkeypatch.setattr(downloader.bs, "login", lambda: None)
    monkeypatch.setattr(downloader.bs, "logout", fake_logout)
    monkeypatch.setattr(
        downloader.bs,
        "query_history_k_data_plus",
        lambda **kwargs: responses.pop(0),
    )

    assert list(download_daily_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
        ("sh600000", None)
    ]
    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)

    stale_file = tmp_path / "toy" / "month=2024-01" / "stale.pq"
    stale_file.parent.mkdir(parents=True)
    batch_frame.iloc[[0]][DATA_COLUMNS[2:]].to_parquet(stale_file, index=False)

    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 not stale_file.exists()
    assert written[DATA_COLUMNS].columns.tolist() == DATA_COLUMNS
    assert written["symbol_name"].tolist() == ["PF Bank", "PF Bank"]


def test_download_universe_raises_when_all_daily_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_daily_batch",
        lambda symbols, start, end, adjust="qfq": iter([("sh600000", None)]),
    )

    with pytest.raises(RuntimeError, match="No data downloaded"):
        download_universe(
            universe="toy",
            start_date="2024-01-02",
            end_date="2024-01-02",
            output_dir=str(tmp_path),
        )


def test_download_universe_progress_branch_at_100_symbols(tmp_path, monkeypatch):
    symbols = [f"sh6{i:05d}" for i in range(100)]
    batch_frame = pd.DataFrame({
        "symbol": ["sh600000"],
        "date": [pd.Timestamp("2024-01-02")],
        "open": [10.0],
        "high": [11.0],
        "low": [9.0],
        "close": [10.5],
        "preclose": [10.0],
        "volume": [1000.0],
        "amount": [10500.0],
        "vwap": [10.5],
        "turn": [1.0],
        "pctChg": [5.0],
        "tradestatus": [1],
        "isST": [0],
        "peTTM": [8.0],
        "pbMRQ": [1.1],
        "psTTM": [2.1],
        "pcfNcfTTM": [3.1],
    })

    monkeypatch.setattr(
        pipeline_downloader,
        "_resolve_universe",
        lambda universe, max_symbols=0: pd.DataFrame({
            "symbol_id": symbols,
            "symbol_name": symbols,
        }),
    )

    def fake_batch(requested_symbols, start, end, adjust="qfq"):
        assert requested_symbols == symbols
        for symbol in requested_symbols:
            yield symbol, batch_frame.copy()

    monkeypatch.setattr(pipeline_downloader, "download_daily_batch", fake_batch)

    stats = download_universe(
        universe="toy100",
        start_date="2024-01-02",
        end_date="2024-01-02",
        output_dir=str(tmp_path),
        chunk_size=200,
    )

    assert stats["n_symbols"] == 100
    assert stats["n_rows"] == 100