from __future__ import annotations import ast import json import unittest from pathlib import Path from typing import Any from backend.data import build_data_gateway from backend.data.datahub.bridge import DatahubAwareTushareClient, DatahubBridge, looks_like_heaven from backend.data.datahub.client import DatahubClient, DatahubResponse from backend.data.datahub.compare import compare_rows from backend.data.datahub.errors import DatahubError from backend.data.datahub.native import to_canonical_row, to_native_row from backend.data.datahub.settings import DATASETS, DatahubSettings, DatasetFlags ROOT = Path(__file__).resolve().parents[1] TOKEN = "super-secret-datahub-token" LEGACY_DAILY = { "ts_code": "600000.SH", "trade_date": "20240902", "open": 10.11, "high": 10.25, "low": 10.01, "close": 10.20, "pct_chg": 1.2345, "vol": 1000.0, "amount": 2000.0, } HUB_DAILY = { "ts_code": "600000.SH", "trade_date": "20240902", "open": 10.11, "high": 10.25, "low": 10.01, "close": 10.20, "pct_chg": 1.2345, "volume": 100000.0, "amount": 2000000.0, } class FakeLegacy: def __init__(self, rows: list[dict[str, Any]] | Exception | None = None) -> None: self.token = "legacy-token" self.timeout = 30 self.rows = [] if rows is None else rows self.calls: list[tuple[str, dict[str, Any] | None, str]] = [] def query(self, api_name: str, params: dict[str, Any] | None = None, fields: str = "") -> list[dict[str, Any]]: self.calls.append((api_name, params, fields)) if isinstance(self.rows, Exception): raise self.rows return [dict(row) for row in self.rows] class FakeClient(DatahubClient): def __init__(self, error: DatahubError | None = None, response: DatahubResponse | None = None) -> None: super().__init__(DatahubSettings(base_url="http://127.0.0.1:9", token=TOKEN)) self.error = error self.response = response or DatahubResponse( data=[dict(HUB_DAILY)], meta={"tier": "official", "trade_date": "20240902", "stale": False, "staleness_seconds": 0}, ) self.paths: list[str] = [] def get(self, path: str, params: dict[str, Any] | None = None) -> DatahubResponse: self.paths.append(path) if TOKEN in json.dumps(params or {}) or TOKEN in path: raise AssertionError("token leaked into url") if self.error: raise self.error return self.response def flags(**enabled: tuple[bool, bool]) -> DatahubSettings: datasets = {name: DatasetFlags(name) for name in DATASETS} for name, pair in enabled.items(): datasets[name] = DatasetFlags(name, read=pair[0], shadow=pair[1]) return DatahubSettings(base_url="http://127.0.0.1:9", token=TOKEN, datasets=datasets) class DatahubBridgeTests(unittest.TestCase): def test_default_config_keeps_legacy_and_does_not_call_datahub(self) -> None: settings = DatahubSettings.load(environ={}, credentials={}) self.assertFalse(settings.any_enabled()) self.assertTrue(all(not settings.flags(name).read and not settings.flags(name).shadow for name in DATASETS)) client = FakeClient(error=DatahubError("INTERNAL", "should not be called")) legacy = FakeLegacy([LEGACY_DAILY]) wrapped = DatahubAwareTushareClient(legacy, DatahubBridge(settings, client)) rows = wrapped.query("daily", {"trade_date": "20240902"}, "ts_code,close,vol,amount") self.assertEqual(rows[0]["amount"], 2000.0) self.assertEqual(client.paths, []) self.assertEqual(len(legacy.calls), 1) def test_each_dataset_has_independent_read_flag(self) -> None: settings = flags(daily=(True, False), auction=(False, False)) self.assertTrue(settings.flags("daily").read) self.assertFalse(settings.flags("auction").read) self.assertFalse(any(settings.flags(name).read for name in DATASETS if name != "daily")) source = (ROOT / "config" / "datahub.config.json").read_text(encoding="utf-8") self.assertNotIn("master", source) self.assertNotIn("DATAHUB_READ_ALL", source) def test_read_flag_replaces_only_that_dataset_and_converts_units(self) -> None: shadows: list[dict[str, Any]] = [] client = FakeClient() legacy = FakeLegacy([LEGACY_DAILY]) wrapped = DatahubAwareTushareClient( legacy, DatahubBridge(flags(daily=(True, False)), client, shadow_sink=shadows.append), ) rows = wrapped.query("daily", {"trade_date": "20240902"}, "ts_code,vol,amount") self.assertEqual(rows[0]["vol"], 1000.0) self.assertEqual(rows[0]["amount"], 2000.0) self.assertEqual(legacy.calls, []) self.assertEqual(client.paths, ["/v1/bars/daily"]) calendar_legacy = FakeLegacy([{"cal_date": "20240902", "is_open": 1}]) calendar_client = FakeClient(error=DatahubError("INTERNAL", "nope")) calendar_wrapped = DatahubAwareTushareClient( calendar_legacy, DatahubBridge(flags(daily=(True, False)), calendar_client), ) calendar = calendar_wrapped.query("trade_cal", {"start_date": "20240902", "end_date": "20240902"}, "") self.assertEqual(calendar[0]["is_open"], 1) self.assertEqual(calendar_client.paths, []) def test_fallback_on_down_401_timeout_empty_unpublished_and_stale(self) -> None: cases = [ DatahubError("UNAVAILABLE", "down"), DatahubError("UNAUTHORIZED", "401"), DatahubError("TIMEOUT", "late"), DatahubError("EMPTY", "no rows"), DatahubError("DATASET_NOT_PUBLISHED", "not ready"), DatahubError("STALE", "old"), ] for error in cases: with self.subTest(error=error.code): if error.code == "EMPTY": client = FakeClient(response=DatahubResponse(data=[], meta={"stale": False, "staleness_seconds": 0})) elif error.code == "STALE": client = FakeClient(response=DatahubResponse( data=[dict(HUB_DAILY)], meta={"stale": True, "staleness_seconds": 999999}, )) else: client = FakeClient(error=error) legacy = FakeLegacy([LEGACY_DAILY]) wrapped = DatahubAwareTushareClient(legacy, DatahubBridge(flags(daily=(True, False)), client)) rows = wrapped.query("daily", {"trade_date": "20240902"}, "ts_code,amount") self.assertEqual(rows[0]["amount"], 2000.0) self.assertEqual(len(legacy.calls), 1) def test_shadow_compares_without_replacing_and_survives_hub_failure(self) -> None: reports: list[dict[str, Any]] = [] client = FakeClient() legacy = FakeLegacy([LEGACY_DAILY]) wrapped = DatahubAwareTushareClient( legacy, DatahubBridge(flags(daily=(False, True)), client, shadow_sink=reports.append), ) rows = wrapped.query("daily", {"trade_date": "20240902"}, "ts_code,amount,vol") self.assertEqual(rows[0]["amount"], 2000.0) self.assertEqual(len(legacy.calls), 1) self.assertEqual(reports[0]["equal"], True) self.assertEqual(reports[0]["matched"], 1) failed = FakeClient(error=DatahubError("UNAVAILABLE", TOKEN)) fail_reports: list[dict[str, Any]] = [] fail_legacy = FakeLegacy([LEGACY_DAILY]) fail_wrapped = DatahubAwareTushareClient( fail_legacy, DatahubBridge(flags(daily=(False, True)), failed, shadow_sink=fail_reports.append), ) again = fail_wrapped.query("daily", {"trade_date": "20240902"}, "amount") self.assertEqual(again[0]["amount"], 2000.0) self.assertTrue(fail_reports[0]["hub_error"]) self.assertNotIn(TOKEN, json.dumps(fail_reports[0])) def test_compare_classifies_unit_conversion_missing_row_and_value_diff(self) -> None: equal = compare_rows("daily", [LEGACY_DAILY], [HUB_DAILY], {"stale": False, "staleness_seconds": 0}) self.assertTrue(equal["equal"]) unit = compare_rows("daily", [LEGACY_DAILY], [{**HUB_DAILY, "amount": 2000.0, "volume": 1000.0}]) self.assertGreater(unit["unit_conversion_count"], 0) missing = compare_rows("daily", [LEGACY_DAILY], []) self.assertEqual(missing["missing_hub_count"], 1) value = compare_rows("daily", [LEGACY_DAILY], [{**HUB_DAILY, "close": 99.0}]) self.assertEqual(value["value_diff_count"], 1) skew = compare_rows("daily", [LEGACY_DAILY], [HUB_DAILY], {"stale": False, "staleness_seconds": 12}) self.assertTrue(skew["time_skew"]) def test_native_roundtrip_matches_known_scales(self) -> None: native = to_native_row("daily", HUB_DAILY) self.assertEqual(native["vol"], 1000.0) self.assertEqual(native["amount"], 2000.0) canonical = to_canonical_row("daily", native) self.assertEqual(canonical["vol"], 100000.0) self.assertEqual(canonical["amount"], 2000000.0) def test_heaven_keeps_legacy_even_when_read_flag_is_on(self) -> None: self.assertTrue(looks_like_heaven("backend.features.heaven.market_context", "backend/features/heaven/market_context.py")) self.assertFalse(looks_like_heaven("backend.features.market.service", "backend/features/market/service.py")) client = FakeClient() legacy = FakeLegacy([LEGACY_DAILY]) wrapped = DatahubAwareTushareClient( legacy, DatahubBridge(flags(daily=(True, False)), client, heaven_guard=lambda: True), ) rows = wrapped.query("daily", {"trade_date": "20240902"}, "amount") self.assertEqual(rows[0]["amount"], 2000.0) self.assertEqual(client.paths, []) def test_status_flag_does_not_run_when_off_and_falls_back_when_on(self) -> None: off = DatahubBridge(flags(), FakeClient(error=DatahubError("UNAVAILABLE", "down"))) self.assertIsNone(off.dataset_status("20240902")) reports: list[dict[str, Any]] = [] failed = DatahubBridge( flags(status=(True, True)), FakeClient(error=DatahubError("UNAUTHORIZED", "nope")), shadow_sink=reports.append, ) self.assertIsNone(failed.dataset_status("20240902")) self.assertTrue(reports[0]["hub_error"]) ok = DatahubBridge( flags(status=(True, False)), FakeClient(response=DatahubResponse(data=[{"dataset": "daily", "state": "published"}], meta={"stale": False, "staleness_seconds": 0})), ) self.assertEqual(ok.dataset_status("20240902")[0]["state"], "published") def test_default_gateway_wraps_tushare_without_calling_datahub(self) -> None: gateway = build_data_gateway({}, datahub_settings=flags()) client = gateway.tushare() self.assertIsInstance(client, DatahubAwareTushareClient) self.assertFalse(gateway.datahub.settings.any_enabled()) def test_features_do_not_import_datahub_client(self) -> None: violations = [] for path in (ROOT / "backend" / "features").rglob("*.py"): tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) for node in ast.walk(tree): names = [] if isinstance(node, ast.Import): names = [alias.name for alias in node.names] elif isinstance(node, ast.ImportFrom) and node.module: names = [node.module] for name in names: if "datahub" in name.split("."): violations.append(f"{path.relative_to(ROOT)} -> {name}") self.assertEqual(violations, []) if __name__ == "__main__": unittest.main()