xiaobai-review/xiaobai-datahub/tests/test_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.error import HTTPError
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, fake_transport

ERROR_CODES = {
    "UNAUTHORIZED",
    "INVALID_ARGUMENT",
    "RATE_LIMITED",
    "SOURCE_UNAVAILABLE",
    "DATASET_NOT_PUBLISHED",
    "STALE_DATA",
    "INTERNAL",
}


class ApiContractTests(unittest.TestCase):
    def setUp(self) -> None:
        self.tmp = tempfile.TemporaryDirectory()
        key = SecretVault.generate_key()
        self.token = "k" * 32
        settings = Settings(
            host="127.0.0.1",
            port=0,
            encryption_key=key,
            api_token=self.token,
            admin_password="StartPass1",
            tushare_token="tushare-secret-token-xyz",
            db_path=Path(self.tmp.name) / "hub.db",
            backup_dir=Path(self.tmp.name) / "backups",
            scheduler_enabled=False,
            quality={"daily_row_ratio": 0.5, "null_rate_max": 0.5, "list_limit_default": 5000, "list_limit_max": 5000},
        )
        adapter = TushareAdapter("tushare-secret-token-xyz", transport=fake_transport)
        self.hub = Hub(settings, adapter=adapter)
        self.hub.pipeline.ingest_reference(TRADE_DATE)
        for dataset in ("daily", "valuation", "moneyflow", "auction", "index_daily"):
            self.hub.pipeline.run_dataset(dataset, TRADE_DATE)
        handler = make_handler(self.hub)
        self.server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self.base = f"http://127.0.0.1:{self.server.server_address[1]}"

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

    def _get(self, path: str, token: str | None = None) -> tuple[int, dict]:
        headers = {}
        if token is not None:
            headers["X-Datahub-Token"] = token
        req = Request(self.base + path, headers=headers)
        try:
            with urlopen(req, timeout=5) as resp:
                return resp.status, json.loads(resp.read().decode())
        except HTTPError as exc:
            return exc.code, json.loads(exc.read().decode())

    def test_livez_no_token(self) -> None:
        status, body = self._get("/livez", token=None)
        self.assertEqual(status, 200)
        self.assertEqual(body["status"], "ok")

    def test_missing_and_bad_token_401(self) -> None:
        status, body = self._get("/v1/health", token=None)
        self.assertEqual(status, 401)
        self.assertEqual(body["error"]["code"], "UNAUTHORIZED")
        status, body = self._get("/v1/health", token="wrong")
        self.assertEqual(status, 401)
        self.assertNotIn("tushare-secret-token-xyz", json.dumps(body))
        self.assertNotIn(self.token, json.dumps(body))

    def test_core_endpoints_schema(self) -> None:
        paths = [
            "/v1/health",
            f"/v1/calendar?from=20240901&to=20240907",
            "/v1/stocks",
            f"/v1/bars/daily?date={TRADE_DATE}&code=600000.SH&adjust=none",
            f"/v1/bars/daily?date={TRADE_DATE}&code=600000.SH&adjust=qfq",
            f"/v1/indexes/bars?date={TRADE_DATE}&code=000001.SH",
            f"/v1/valuation?date={TRADE_DATE}&code=600000.SH",
            f"/v1/moneyflow?date={TRADE_DATE}&code=600000.SH",
            f"/v1/auction?date={TRADE_DATE}",
            f"/v1/datasets/status?date={TRADE_DATE}",
            f"/v1/batches?date={TRADE_DATE}",
        ]
        for path in paths:
            status, body = self._get(path, token=self.token)
            self.assertEqual(status, 200, path)
            self.assertEqual(body["schema_version"], 1)
            self.assertIn("data", body)
            self.assertIn("meta", body)
            self.assertIn("tier", body["meta"])

    def test_qfq_matches_formula(self) -> None:
        _, none = self._get(f"/v1/bars/daily?date={TRADE_DATE}&code=600000.SH&adjust=none", token=self.token)
        _, qfq = self._get(f"/v1/bars/daily?date={TRADE_DATE}&code=600000.SH&adjust=qfq", token=self.token)
        raw = none["data"][0]
        adj = qfq["data"][0]
        expected = round(raw["close"] * raw["adj_factor"] / raw["adj_factor"], 4)
        self.assertEqual(adj["close"], expected)

    def test_unpublished_code(self) -> None:
        status, body = self._get("/v1/bars/daily?date=19990101", token=self.token)
        self.assertEqual(status, 404)
        self.assertEqual(body["error"]["code"], "DATASET_NOT_PUBLISHED")
        self.assertIn("expected_at", body["error"])

    def test_error_code_set_documented(self) -> None:
        self.assertGreaterEqual(ERROR_CODES, {"UNAUTHORIZED", "DATASET_NOT_PUBLISHED", "INVALID_ARGUMENT"})

    def test_six_digit_code(self) -> None:
        status, body = self._get(f"/v1/bars/daily?date={TRADE_DATE}&code=600000", token=self.token)
        self.assertEqual(status, 200)
        self.assertEqual(body["data"][0]["ts_code"], "600000.SH")

    def test_amount_unit_is_yuan(self) -> None:
        _, body = self._get(f"/v1/bars/daily?date={TRADE_DATE}&code=600000.SH", token=self.token)
        self.assertEqual(body["data"][0]["amount"], 2_000_000.0)
        _, flow = self._get(f"/v1/moneyflow?date={TRADE_DATE}&code=600000.SH", token=self.token)
        self.assertEqual(flow["data"][0]["net_mf_amount"], 170000.0)

    def test_token_never_in_health_or_admin_sources(self) -> None:
        _, health = self._get("/v1/health", token=self.token)
        blob = json.dumps(health)
        self.assertNotIn("tushare-secret-token-xyz", blob)
        self.assertNotIn(self.token, blob)


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