xiaobai-review/xiaobai-datahub/tests/test_admin_observability_api.py

from __future__ import annotations

import json
import tempfile
import threading
import unittest
from http.server import ThreadingHTTPServer
from pathlib import Path
from urllib.request import Request, urlopen

from datahub.adapters.tushare import TushareAdapter
from datahub.crypto import SecretVault
from datahub.httpapp import make_handler
from datahub.hub import Hub
from datahub.settings import Settings
from tests.fixtures import TRADE_DATE, StubSiteAuth, fake_transport


class AdminObservabilityApiTests(unittest.TestCase):
    """HEL-543: new read-only admin endpoints for provider status, source
    catalog and lineage. These must never require write access and must
    never touch the existing routing/publish logic."""

    def setUp(self) -> None:
        self.tmp = tempfile.TemporaryDirectory()
        settings = Settings(
            encryption_key=SecretVault.generate_key(),
            api_token="z" * 32,
            tushare_token="real-tushare-token-abcdef",
            db_path=Path(self.tmp.name) / "hub.db",
            scheduler_enabled=False,
        )
        self.site_auth = StubSiteAuth()
        self.hub = Hub(
            settings,
            adapter=TushareAdapter("real-tushare-token-abcdef", transport=fake_transport),
            site_auth=self.site_auth,
        )
        handler = make_handler(self.hub)
        self.server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
        threading.Thread(target=self.server.serve_forever, daemon=True).start()
        self.base = f"http://127.0.0.1:{self.server.server_address[1]}"
        self.cookie = f"xiaobai_session={StubSiteAuth.SESSION}"
        self.csrf = self.site_auth.csrf_token(StubSiteAuth.SESSION)

    def tearDown(self) -> None:
        self.server.shutdown()
        self.server.server_close()
        self.tmp.cleanup()

    def _json(self, path, method="GET", body=None, cookie="", csrf=""):
        data = None if body is None else json.dumps(body).encode()
        headers = {"Content-Type": "application/json"}
        if cookie:
            headers["Cookie"] = cookie
        if csrf:
            headers["X-CSRF-Token"] = csrf
        req = Request(self.base + path, data=data, headers=headers, method=method)
        with urlopen(req, timeout=5) as resp:
            set_cookie = resp.headers.get("Set-Cookie", "")
            return resp.status, json.loads(resp.read().decode()), set_cookie

    def _get(self, path):
        return self._json(path, cookie=self.cookie, csrf=self.csrf)

    def test_providers_status_reflects_real_pipeline_activity(self) -> None:
        pipeline = self.hub.pipeline
        pipeline.ingest_reference(TRADE_DATE)
        pipeline.run_dataset("daily", TRADE_DATE)

        status, body, _ = self._get("/admin/api/providers/status")
        self.assertEqual(status, 200)
        health = body["health"]
        self.assertTrue(any(item["provider"] == "tushare" and item["interface"] == "daily" for item in health))
        row = next(item for item in health if item["provider"] == "tushare" and item["interface"] == "daily")
        self.assertEqual(row["state"], "ok")
        recent = body["recent_calls"]
        self.assertTrue(any(item["provider"] == "tushare" and item["interface"] == "daily" for item in recent))

    def test_providers_status_filters_by_provider(self) -> None:
        pipeline = self.hub.pipeline
        pipeline.ingest_reference(TRADE_DATE)
        pipeline.run_dataset("daily", TRADE_DATE)

        status, body, _ = self._get("/admin/api/providers/status?provider=tushare")
        self.assertEqual(status, 200)
        self.assertTrue(all(item["provider"] == "tushare" for item in body["health"]))
        self.assertTrue(all(item["provider"] == "tushare" for item in body["recent_calls"]))

        status, body, _ = self._get("/admin/api/providers/status?provider=eastmoney")
        self.assertEqual(status, 200)
        self.assertEqual(body["health"], [])
        self.assertEqual(body["recent_calls"], [])

    def test_source_catalog_lists_known_providers_without_leaking_secrets(self) -> None:
        status, body, _ = self._get("/admin/api/source-catalog")
        self.assertEqual(status, 200)
        blob = json.dumps(body)
        self.assertNotIn("real-tushare-token-abcdef", blob)
        providers = {item["provider"] for item in body["items"]}
        self.assertIn("tushare", providers)
        self.assertIn("eastmoney", providers)
        self.assertIn("tencent", providers)
        self.assertIn("ifind", providers)

    def test_lineage_snapshot_and_affected_query(self) -> None:
        status, body, _ = self._get("/admin/api/lineage")
        self.assertEqual(status, 200)
        self.assertTrue(len(body["items"]) > 0)
        datasets = {item["dataset"] for item in body["items"]}
        self.assertIn("stocks", datasets)

        status, body, _ = self._get("/admin/api/lineage/affected?provider=tushare&interface=daily")
        self.assertEqual(status, 200)
        self.assertEqual(body["provider"], "tushare")
        self.assertEqual(body["interface"], "daily")

    def test_disabled_kill_switch_reports_enabled_false_with_empty_structure(self) -> None:
        # HEL-543 total-review 🔴: flip the runtime kill switch the same way
        # Hub.__init__ wires Settings.observability_enabled onto the db
        # handle, then confirm every new endpoint reports disabled with an
        # explicit empty structure rather than silently going quiet.
        self.hub.db.observability_enabled = False
        try:
            pipeline = self.hub.pipeline
            pipeline.ingest_reference(TRADE_DATE)
            pipeline.run_dataset("daily", TRADE_DATE)  # must still fully succeed

            status, body, _ = self._get("/admin/api/providers/status")
            self.assertEqual(status, 200)
            self.assertEqual(body, {"enabled": False, "health": [], "recent_calls": []})

            status, body, _ = self._get("/admin/api/source-catalog")
            self.assertEqual(status, 200)
            self.assertEqual(body, {"enabled": False, "items": []})

            status, body, _ = self._get("/admin/api/lineage")
            self.assertEqual(status, 200)
            self.assertFalse(body["enabled"])
            self.assertEqual(body["items"], [])

            status, body, _ = self._get("/admin/api/lineage/affected?provider=tushare")
            self.assertEqual(status, 200)
            self.assertEqual(
                body, {"enabled": False, "provider": "tushare", "interface": "", "items": []}
            )

            # Nothing was ever written while disabled.
            self.assertEqual(self.hub.db.fetchall("SELECT * FROM provider_call_log"), [])
        finally:
            self.hub.db.observability_enabled = True

    def test_non_admin_site_accounts_cannot_read_the_new_endpoints(self) -> None:
        from urllib.error import HTTPError

        self.site_auth.add_session("member-session", StubSiteAuth.MEMBER)
        with self.assertRaises(HTTPError) as ctx:
            self._json("/admin/api/source-catalog", cookie="xiaobai_session=member-session")
        self.assertEqual(ctx.exception.code, 403)

        with self.assertRaises(HTTPError) as ctx:
            self._json("/admin/api/source-catalog")
        self.assertEqual(ctx.exception.code, 401)


if __name__ == "__main__":
    unittest.main()