zhixing-system/zhixing-server/tests/unit/market_data/test_tushare.py

from datetime import date

import pytest
import tushare as ts  # pyright: ignore[reportMissingTypeStubs]

from zhixing_server.modules.market_data.domain.models import SyncWindow
from zhixing_server.modules.market_data.infrastructure.tushare import (
    RequestCoordinator,
    TushareAdapter,
)


def test_from_token_reuses_api_client_for_pro_bar(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    created_client = object()
    calls: list[dict[str, object]] = []

    def fake_pro_api(token: str) -> object:
        assert token
        return created_client

    def fake_pro_bar(**kwargs: object) -> list[dict[str, object]]:
        calls.append(kwargs)
        return [
            {
                "ts_code": "000001.SZ",
                "trade_date": "20240102",
                "close": "10",
            }
        ]

    monkeypatch.setattr(ts, "pro_api", fake_pro_api)
    monkeypatch.setattr(ts, "pro_bar", fake_pro_bar)

    adapter = TushareAdapter.from_token(
        "test-token",
        request_interval_seconds=0,
    )

    bars = adapter.fetch_bars(
        "000001.SZ",
        SyncWindow(start=date(2024, 1, 2), end=date(2024, 1, 2)),
    )

    assert bars[0].ts_code == "000001.SZ"
    assert calls
    assert calls[0]["api"] is created_client
    assert calls[0]["adj"] == "qfq"
    assert calls[0]["retry_count"] == 1


def test_rate_limit_cooldown_is_shared_by_following_requests() -> None:
    current = [0.0]
    waits: list[float] = []
    calls: list[float] = []

    def clock() -> float:
        return current[0]

    def wait(seconds: float) -> None:
        waits.append(seconds)
        current[0] += seconds

    coordinator = RequestCoordinator(
        max_retries=0,
        cooldown_seconds=(60, 120, 180),
        clock=clock,
        wait_fn=wait,
        sleep_fn=wait,
    )

    def rate_limited() -> object:
        calls.append(current[0])
        raise RuntimeError("HTTP 429: too many requests")

    with pytest.raises(RuntimeError):
        coordinator.call("daily", rate_limited)

    assert calls == [0]
    assert coordinator.cooldown_until == 60

    result = coordinator.call("adj_factor", lambda: calls.append(current[0]) or "ok")

    assert result == "ok"
    assert calls == [0, 60]
    assert waits == [60]


def test_pro_bar_qfq_calls_are_bound_to_the_shared_coordinator(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class Client:
        def __init__(self) -> None:
            self.calls: list[str] = []

        def daily(self, **kwargs: object) -> object:
            self.calls.append("daily")
            return object()

        def adj_factor(self, **kwargs: object) -> object:
            self.calls.append("adj_factor")
            return object()

    client = Client()
    pro_bar_calls: list[dict[str, object]] = []

    def fake_pro_api(token: str) -> object:
        return client

    def fake_pro_bar(**kwargs: object) -> list[dict[str, object]]:
        pro_bar_calls.append(kwargs)
        api = kwargs["api"]
        assert api is client
        api.daily(ts_code="000001.SZ")  # type: ignore[union-attr]
        api.adj_factor(ts_code="000001.SZ")  # type: ignore[union-attr]
        return [{"ts_code": "000001.SZ", "trade_date": "20240102", "close": "10"}]

    monkeypatch.setattr(ts, "pro_api", fake_pro_api)
    monkeypatch.setattr(ts, "pro_bar", fake_pro_bar)

    adapter = TushareAdapter.from_token("test-token", request_interval_seconds=0)
    adapter.fetch_bars(
        "000001.SZ",
        SyncWindow(start=date(2024, 1, 2), end=date(2024, 1, 2)),
    )

    assert client.calls == ["daily", "adj_factor"]
    assert pro_bar_calls[0]["retry_count"] == 1