from __future__ import annotations import re import unittest from http import HTTPStatus from pathlib import Path from typing import Any ROOT = Path(__file__).resolve().parents[1] REVIEW_ROOT = ROOT.parent from datahub.serving import ApiError from datahub.siteauth import SiteBridgeError from datahub.siteauth import SiteBridgeError from datahub.siteconsole import SiteConsole class RecordingBridge: """Stands in for the review site so these tests exercise only the mapping.""" def __init__(self, replies: dict[str, dict[str, Any]] | None = None) -> None: self.replies = replies or {} self.calls: list[tuple[str, dict[str, Any]]] = [] self.fail_with = "" def call(self, path: str, payload: dict[str, Any] | None = None) -> dict[str, Any]: self.calls.append((path, dict(payload or {}))) if self.fail_with: raise SiteBridgeError(self.fail_with) return self.replies.get(path, {}) def paths(self) -> list[str]: return [path for path, _ in self.calls] STATUS = { "llm": { "primary_model_id": "gpt-main", "fallback_model_id": "", "models": [ {"id": "gpt-main", "name": "GPT 主力", "model": "gpt-4o", "base_url": "https://api.openai.com/v1", "configured": True, "api_key_last4": "7f21"}, {"id": "gpt-mini", "name": "GPT 轻量", "model": "gpt-4o-mini", "base_url": "https://api.openai.com/v1", "configured": True, "api_key_last4": "7f21"}, {"id": "self-host", "name": "自建 Qwen", "model": "qwen2.5", "base_url": "https://llm.intra.example.com/v1", "configured": False, "api_key_last4": ""}, ], } } class ModelPoolTests(unittest.TestCase): def setUp(self) -> None: self.bridge = RecordingBridge({"/api/hub-admin/status": STATUS}) self.console = SiteConsole(self.bridge) def test_models_are_grouped_per_vendor_so_one_vendor_can_hold_many(self) -> None: payload = self.console.models() groups = {group["base_url"]: group for group in payload["groups"]} self.assertEqual(2, len(groups)) openai = groups["https://api.openai.com/v1"] self.assertEqual("OpenAI", openai["label"]) self.assertEqual(2, len(openai["models"])) self.assertEqual("7f21", openai["key_last4"]) self.assertTrue(openai["configured"]) def test_unknown_vendors_fall_back_to_their_host_as_a_label(self) -> None: groups = {group["base_url"]: group for group in self.console.models()["groups"]} self.assertEqual("llm.intra.example.com", groups["https://llm.intra.example.com/v1"]["label"]) self.assertFalse(groups["https://llm.intra.example.com/v1"]["configured"]) def test_vendor_presets_are_offered_for_the_new_vendor_picker(self) -> None: vendors = self.console.models()["vendors"] self.assertIn("OpenAI", [vendor["label"] for vendor in vendors]) self.assertTrue(all(vendor["base_url"].startswith("http") for vendor in vendors)) def test_saving_only_forwards_the_keys_the_caller_actually_sent(self) -> None: self.console.save_models({"primary_model_id": "gpt-mini"}) path, payload = self.bridge.calls[0] self.assertEqual("/api/hub-admin/settings/save", path) self.assertEqual({"primary_model_id": "gpt-mini"}, payload) def test_saving_nothing_is_refused_rather_than_wiping_the_pool(self) -> None: with self.assertRaises(ApiError): self.console.save_models({}) self.assertEqual([], self.bridge.paths()) def test_fetching_a_model_list_needs_a_vendor_endpoint(self) -> None: with self.assertRaises(ApiError): self.console.fetch_models({"api_key": "sk-test"}) def test_fetch_passes_the_typed_key_through_for_first_time_vendors(self) -> None: self.bridge.replies["/api/hub-admin/models/fetch"] = {"models": ["gpt-4o", "gpt-4o-mini"]} result = self.console.fetch_models({"base_url": "https://api.openai.com/v1", "api_key": "sk-new"}) self.assertEqual(["gpt-4o", "gpt-4o-mini"], result["models"]) self.assertEqual({"base_url": "https://api.openai.com/v1", "api_key": "sk-new"}, self.bridge.calls[0][1]) def test_a_site_outage_becomes_a_console_error_not_a_traceback(self) -> None: self.bridge.fail_with = "主站不可达" with self.assertRaises(ApiError) as caught: self.console.models() self.assertIn("主站不可达", str(caught.exception)) class MemberAndQuotaTests(unittest.TestCase): def setUp(self) -> None: self.bridge = RecordingBridge({"/api/hub-admin/members": {"users": [{"id": 2}], "membership": {"member_daily_limit": 50}}}) self.console = SiteConsole(self.bridge) def test_saving_a_member_reads_the_roster_back_so_the_page_shows_truth(self) -> None: payload = self.console.save_member({"user_id": 2, "status": "active", "duration": "3_months"}) self.assertEqual(["/api/hub-admin/membership/save", "/api/hub-admin/members"], self.bridge.paths()) self.assertEqual([{"id": 2}], payload["users"]) def test_quota_is_a_system_setting_not_a_membership_row(self) -> None: self.console.save_quota({"member_daily_limit": 80}) self.assertEqual("/api/hub-admin/settings/save", self.bridge.paths()[0]) self.assertEqual({"member_daily_limit": 80}, self.bridge.calls[0][1]) def test_quota_below_one_is_rejected_before_it_reaches_the_site(self) -> None: for bad in (0, -5, "abc"): with self.subTest(bad=bad): with self.assertRaises(ApiError): self.console.save_quota({"member_daily_limit": bad}) self.assertEqual([], self.bridge.paths()) class InviteTests(unittest.TestCase): def setUp(self) -> None: self.bridge = RecordingBridge( { "/api/hub-admin/invites": {"summary": {"unused": 1}, "codes": [{"code_id": "abc", "code_masked": "XB-9Q2F-••••"}]}, "/api/hub-admin/invites/create": { "created": [{"code_id": "abc", "code": "XB-9Q2F-7K3M-2P8T"}], "summary": {"unused": 1}, "codes": [{"code_id": "abc", "code_masked": "XB-9Q2F-••••"}], }, } ) self.console = SiteConsole(self.bridge) def test_plaintext_codes_come_back_only_from_the_create_call(self) -> None: created = self.console.create_invites({"count": 1}, created_by=1) self.assertEqual("XB-9Q2F-7K3M-2P8T", created["created"][0]["code"]) # 列表里永远只有掩码，完整码不会再出现第二次 listed = self.console.invites() self.assertEqual("XB-9Q2F-••••", listed["codes"][0]["code_masked"]) self.assertNotIn("code", listed["codes"][0]) def test_the_operator_is_recorded_as_the_issuer(self) -> None: self.console.create_invites({"count": 3, "note": "给张总"}, created_by=7) payload = self.bridge.calls[0][1] self.assertEqual(7, payload["created_by"]) self.assertEqual(3, payload["count"]) self.assertEqual("给张总", payload["note"]) def test_revoking_uses_the_public_handle_never_the_raw_code(self) -> None: self.console.revoke_invite({"code_id": "abc"}) self.assertEqual("/api/hub-admin/invites/revoke", self.bridge.paths()[0]) self.assertEqual({"code_id": "abc"}, self.bridge.calls[0][1]) def test_revoking_without_a_target_is_refused(self) -> None: with self.assertRaises(ApiError): self.console.revoke_invite({}) self.assertEqual([], self.bridge.paths()) if __name__ == "__main__": unittest.main() class BridgeContractTests(unittest.TestCase): """Both halves of the bridge must agree on the path spelling. A typo here fails only at runtime with a confusing 401 (the review site falls through to its browser-session guard), so it is worth a static check. """ def test_every_path_the_console_calls_is_registered_on_the_review_site(self) -> None: console_paths = set() for module in ("siteauth.py", "siteconsole.py"): source = (ROOT / "datahub" / module).read_text(encoding="utf-8") console_paths.update(re.findall(r'"(/api/hub-admin/[a-z/-]+)"', source)) self.assertTrue(console_paths) registry = (REVIEW_ROOT / "backend" / "http" / "dispatch.py").read_text(encoding="utf-8") registered = set(re.findall(r'"(/api/hub-admin/[a-z/-]+)":', registry)) self.assertEqual( set(), console_paths - registered, "控制台在调、主站没注册的桥接路径会静默变成 401", ) class BridgeErrorMappingTests(unittest.TestCase): """主站拒绝入参（Key 不对）不能在中枢这边冒成 500。""" def _console(self, error: SiteBridgeError) -> SiteConsole: class Failing: def call(self, path, payload=None): raise error return SiteConsole(Failing()) def test_upstream_rejection_comes_back_as_a_bad_request(self) -> None: console = self._console(SiteBridgeError("模型列表拉取失败（HTTP 401）", 400)) with self.assertRaises(ApiError) as caught: console.fetch_models({"base_url": "https://api.openai.com/v1", "api_key": "sk-bad"}) self.assertEqual(caught.exception.code, "INVALID_ARGUMENT") self.assertEqual(caught.exception.status, HTTPStatus.BAD_REQUEST) def test_unreachable_site_comes_back_as_service_unavailable(self) -> None: console = self._console(SiteBridgeError("主站不可达：connection refused")) with self.assertRaises(ApiError) as caught: console.models() self.assertEqual(caught.exception.code, "SOURCE_UNAVAILABLE") self.assertEqual(caught.exception.status, HTTPStatus.SERVICE_UNAVAILABLE)