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