Files
xiaobai-review/xiaobai-datahub/tests/test_api.py
T
2026-09-02 22:16:05 +08:00

151 lines
6.2 KiB
Python

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"])
if "calendar" in path or "bars" in path or "indexes" in path or "valuation" in path or "moneyflow" in path or "auction" in path:
self.assertIn("coverage", body["meta"])
self.assertIn("incomplete", 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()