xiaobai-review/tests/test_datahub_bridge.py

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_stale_and_incomplete(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"),
            DatahubError("INCOMPLETE", "truncated"),
        ]
        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},
                    ))
                elif error.code == "INCOMPLETE":
                    client = FakeClient(response=DatahubResponse(
                        data=[dict(HUB_DAILY)],
                        meta={
                            "stale": False,
                            "staleness_seconds": 0,
                            "incomplete": True,
                            "coverage": {"complete": False, "missing_count": 80},
                        },
                    ))
                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_on_first_batch_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_stock_detail_range_query_is_not_silently_accepted_when_incomplete(self) -> None:
        source = (ROOT / "backend" / "data" / "providers" / "tushare_stocks.py").read_text(encoding="utf-8")
        self.assertIn('"daily"', source)
        self.assertIn("start_date", source)
        self.assertIn("end_date", source)
        client = FakeClient(
            response=DatahubResponse(
                data=[dict(HUB_DAILY)],
                meta={"stale": False, "staleness_seconds": 0, "incomplete": True, "coverage": {"complete": False, "missing_count": 89}},
            )
        )
        legacy = FakeLegacy([LEGACY_DAILY])
        wrapped = DatahubAwareTushareClient(legacy, DatahubBridge(flags(daily=(True, False)), client))
        rows = wrapped.query(
            "daily",
            {"ts_code": "600000.SH", "start_date": "20240301", "end_date": "20240902"},
            "ts_code,amount",
        )
        self.assertEqual(rows[0]["amount"], 2000.0)
        self.assertEqual(len(legacy.calls), 1)

    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()