PolyGateway/tests/unit/test_client.py

"""GatewayClient 装配与端到端(fake 后端 + MockTransport)测试(设计 §2.4)。"""

import asyncio
import json
import sys
from pathlib import Path

import httpx
import pytest

from polygateway import (
    AllSourcesExhausted,
    GatewayClient,
    GatewaySettings,
    gather_bounded,
)
from polygateway.backends.memory.breaker import InMemoryGate
from polygateway.backends.memory.cache import InMemoryCache
from polygateway.backends.memory.limiter import InMemoryLimiter
from polygateway.sources import RoundRobinSelector
from polygateway.structured.json_repair import JsonRepairStrategy
from polygateway.structured.native_schema import NativeSchemaStrategy
from polygateway.transports.openai_compat import OpenAICompatTransport
from polygateway.types import (
    BackpressurePolicy,
    BreakerConfig,
    GlobalLimits,
    RetryPolicy,
    SourceConfig,
)

_REPO = Path(__file__).resolve().parents[2]

_ENV = {
    "LLM__QWEN__1__BASE_URL": "https://gw.example/v1",
    "LLM__QWEN__1__API_KEY": "sk-a",
    "LLM__QWEN__1__MODEL": "qwen-max",
    "LLM__QWEN__1__TIMEOUT_S": "120",
    "LLM_MAX_RETRIES": "3",
    "LLM_RETRY_BASE_DELAY": "2.0",
    "LLM_RETRY_MAX_DELAY": "30.0",
    "LLM_CIRCUIT_BREAKER_THRESHOLD": "5",
    "LLM_CIRCUIT_BREAKER_COOLDOWN": "60",
    "PGW_CACHE_BACKEND": "none",
    "PGW_TELEMETRY_BACKEND": "none",
}


def _sse(content='{"answer": 1}'):
    chunk = json.dumps({"choices": [{"delta": {"content": content}}]})
    usage = json.dumps({"choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 4}})
    body = f"data: {chunk}\n\ndata: {usage}\n\ndata: [DONE]\n\n"
    return httpx.Response(200, content=body.encode(), headers={"content-type": "text/event-stream"})


def _source(name="qwen_1", **overrides):
    base = {
        "name": name,
        "provider": "qwen",
        "base_url": "https://gw.example/v1",
        "api_key": "sk",
        "model": "qwen-max",
        "timeout_s": 10.0,
    }
    base.update(overrides)
    return SourceConfig(**base)


def _client(sources=None, handler=None, *, limiter=None, quota_full="wait", **overrides):
    sources = sources or [_source()]
    handler = handler or (lambda request: _sse())
    transport = OpenAICompatTransport(
        client_factory=lambda src: httpx.AsyncClient(transport=httpx.MockTransport(handler))
    )
    defaults = {
        "scope": "llm",
        "sources": sources,
        "selector": RoundRobinSelector(),
        "limiter": limiter
        or InMemoryLimiter(
            scope="llm",
            sources={s.name: s for s in sources},
            global_limits=GlobalLimits(0, 0, 0),
        ),
        "breaker": InMemoryGate(config=BreakerConfig(5, 60.0, 120.0)),
        "transport": transport,
        "retry": RetryPolicy(3, 2.0, 30.0),
        "backpressure": BackpressurePolicy(300.0, 0.01),
        "quota_full": quota_full,
        "structured_strategy": JsonRepairStrategy(),
    }
    defaults.update(overrides)
    return GatewayClient(**defaults)


class TestChatEndToEnd:
    async def test_plain_chat(self):
        async with _client() as client:
            resp = await client.chat([{"role": "user", "content": "hi"}])
        assert resp.content == '{"answer": 1}'
        assert resp.source_name == "qwen_1" and resp.prompt_tokens == 3

    async def test_structured_json_tier(self):
        async with _client() as client:
            resp = await client.chat([{"role": "user", "content": "hi"}], structured="json")
        assert resp.structured_data == {"answer": 1}

    async def test_cache_hit_roundtrip(self):
        client = _client(cache=InMemoryCache(), cache_namespace="proj", cache_ttl_s=3600)
        async with client:
            first = await client.chat([{"role": "user", "content": "hi"}])
            second = await client.chat([{"role": "user", "content": "hi"}])
        assert first.cache_hit is False and second.cache_hit is True

    async def test_structured_unavailable_fails_loudly(self):
        async with _client(structured_strategy=None) as client:
            with pytest.raises(ImportError, match="structured"):
                await client.chat([{"role": "user", "content": "hi"}], structured="json")


class TestSamplingOverlay:
    """调用级采样参数入口(issue #4 Task 3)。"""

    def _capturing_client(self, captured, **overrides):
        def handler(request):
            captured.append(json.loads(request.content))
            return _sse()

        return _client(handler=handler, **overrides)

    async def test_overlay_reaches_request_body(self):
        captured = []
        async with self._capturing_client(captured) as client:
            await client.chat([{"role": "user", "content": "hi"}], overlay={"seed": 42})
        assert captured[0]["seed"] == 42

    async def test_call_level_beats_config_level(self):
        """优先级: 调用级 > 配置级(设计决策 A)。"""
        captured = []
        source = _source(extra_body={"temperature": 0, "top_p": 0.9})
        async with self._capturing_client(captured, sources=[source]) as client:
            await client.chat([{"role": "user", "content": "hi"}], overlay={"temperature": 1})
        assert captured[0]["temperature"] == 1  # 调用级覆盖
        assert captured[0]["top_p"] == 0.9  # 配置级未被顶掉的键保留

    async def test_structured_injection_beats_call_level(self):
        """结构化注入优先级最高: 它关系到响应能否被解析(设计决策 A)。"""
        captured = []
        client = self._capturing_client(captured, structured_strategy=NativeSchemaStrategy())
        async with client:
            await client.chat(
                [{"role": "user", "content": "hi"}],
                structured="json",
                overlay={"response_format": {"type": "text"}},
            )
        assert captured[0]["response_format"] != {"type": "text"}

    async def test_protected_key_rejected_before_onion(self):
        """保护键在进洋葱之前就报错,transport 一次都不该被碰到。"""
        captured = []
        async with self._capturing_client(captured) as client:
            with pytest.raises(ValueError, match="stream"):
                await client.chat([{"role": "user", "content": "hi"}], overlay={"stream": False})
        assert captured == []

    async def test_unserializable_value_rejected_before_onion(self):
        """裸 TypeError 会在 CacheMW 的降级 try 之外炸且无遥测(设计决策 B)。"""
        captured = []
        async with self._capturing_client(captured) as client:
            with pytest.raises(ValueError, match="JSON"):
                await client.chat(
                    [{"role": "user", "content": "hi"}], overlay={"temperature": object()}
                )
        assert captured == []

    async def test_caller_dict_mutation_does_not_leak(self):
        """调用方逐次改 seed 复用同一 dict 是预期模式(设计决策 E)。"""
        captured = []
        caller_overlay = {"seed": 1}
        async with self._capturing_client(captured) as client:
            await client.chat([{"role": "user", "content": "hi"}], overlay=caller_overlay)
            caller_overlay["seed"] = 2
            await client.chat([{"role": "user", "content": "hi"}], overlay=caller_overlay)
        assert [c["seed"] for c in captured] == [1, 2]


class _MemoryRecorder:
    """收下遥测行原样存起来;断言"哪些行被写了"必须能看到零行的情形。"""

    def __init__(self):
        self.rows = []

    async def record_llm_call(self, **fields):
        self.rows.append(fields)


class TestCallerDimensions:
    """调用方自定义维度进遥测(issue #11 Task 4)。"""

    _MSG = [{"role": "user", "content": "hi"}]

    async def test_dimensions_reach_telemetry_row(self):
        recorder = _MemoryRecorder()
        async with _client(telemetry=recorder) as client:
            await client.chat(self._MSG, tenant_id="t1", meta={"batch": "b-42"})
        row = recorder.rows[-1]
        assert row["tenant_id"] == "t1"
        assert json.loads(row["meta"]) == {"batch": "b-42"}

    async def test_default_path_writes_sentinels(self):
        """不传两参数时落哨兵值而非 NULL(§4.4: NULL 在 RLS 下是永久不可见的黑洞)。"""
        recorder = _MemoryRecorder()
        async with _client(telemetry=recorder) as client:
            await client.chat(self._MSG)
        row = recorder.rows[-1]
        assert row["tenant_id"] == "" and row["meta"] == "{}"

    async def test_invalid_meta_key_rejected_before_any_telemetry(self):
        """校验早于遥测(§4.2 核心承诺): 放进洋葱就会被降级成 warning 而调用照常发出。"""
        recorder = _MemoryRecorder()
        async with _client(telemetry=recorder) as client:
            with pytest.raises(ValueError, match="meta"):
                await client.chat(self._MSG, meta={"BAD-KEY": 1})
        assert recorder.rows == []

    async def test_non_finite_float_rejected_before_any_telemetry(self):
        """nan 产出的是 PG 拒收的非法 JSON;放行等于把调用方 bug 变成静默丢遥测(§6)。"""
        recorder = _MemoryRecorder()
        async with _client(telemetry=recorder) as client:
            with pytest.raises(ValueError, match="nan"):
                await client.chat(self._MSG, meta={"k": float("nan")})
        assert recorder.rows == []

    async def test_meta_does_not_enter_cache_key(self):
        """仅 meta 不同必须仍命中缓存(F1): 进 key 会让存量缓存全量冷启动且不报错。

        带对照组: 只断言"命中"的话,缓存 key 退化成常量(忽略一切输入)时本用例
        照样绿——那是恒真断言。故再改一个**确实进 key** 的维度(namespace)断言
        miss,证明 key 仍在区分输入,"meta 不进 key"才是被测出来的结论。
        """
        cache = InMemoryCache()  # 两个 client 共用一份存储,否则对照组的 miss 是白来的
        client = _client(cache=cache, cache_namespace="proj", cache_ttl_s=3600)
        async with client:
            first = await client.chat(self._MSG, meta={"batch": "b-1"})
            second = await client.chat(self._MSG, meta={"batch": "b-2"})
        assert first.cache_hit is False and second.cache_hit is True

        other_ns = _client(cache=cache, cache_namespace="other", cache_ttl_s=3600)
        async with other_ns:
            assert (await other_ns.chat(self._MSG, meta={"batch": "b-1"})).cache_hit is False


class TestModelFingerprint:
    """配置级采样参数须进缓存身份,否则改 temperature 后仍读旧缓存(决策 C)。"""

    def test_empty_extra_body_keeps_legacy_fingerprint(self):
        """全源无 extra_body 时字面量与旧实现逐字相同,不触发存量缓存冷启动。"""
        from polygateway.client import build_model_fingerprint

        sources = [_source(), _source(name="qwen_2", model="qwen-plus")]
        assert build_model_fingerprint(sources) == "qwen-max,qwen-plus"

    def test_extra_body_changes_fingerprint(self):
        from polygateway.client import build_model_fingerprint

        plain = build_model_fingerprint([_source()])
        tuned = build_model_fingerprint([_source(extra_body={"temperature": 0})])
        assert plain != tuned
        assert tuned.startswith("qwen-max|")  # 旧字面量仍是前缀,便于人眼辨认

    def test_enable_thinking_changes_fingerprint(self):
        """issue #5 配套: thinking 一旦真正改变请求体,就必须进缓存身份。

        否则"关掉推理后重启"会读到开着推理时缓存的旧响应——issue #4 为
        temperature 写过逐字相同的理由。
        """
        from polygateway.client import build_model_fingerprint

        plain = build_model_fingerprint([_source()])
        off = build_model_fingerprint([_source(enable_thinking=False)])
        on = build_model_fingerprint([_source(enable_thinking=True)])
        assert len({plain, off, on}) == 3

    def test_extra_body_only_fingerprint_is_byte_identical_to_before(self):
        """只配 extra_body、不表态 thinking 的存量源不得触发冷启动。

        字面量在此硬编码: 这条断言的价值全在"逐字相同",改实现时必须先看见它红。
        """
        import hashlib
        import json

        from polygateway.client import build_model_fingerprint

        mark = json.dumps(["qwen-max", {"temperature": 0}], sort_keys=True, ensure_ascii=False)
        expected = "qwen-max|" + hashlib.sha256(mark.encode("utf-8")).hexdigest()
        assert build_model_fingerprint([_source(extra_body={"temperature": 0})]) == expected

    def test_source_rename_does_not_change_fingerprint(self):
        """指纹按 (model, extra_body) 而非源名: 改名不该误触全量冷启动。"""
        from polygateway.client import build_model_fingerprint

        a = build_model_fingerprint([_source(name="qwen_1", extra_body={"temperature": 0})])
        b = build_model_fingerprint([_source(name="renamed", extra_body={"temperature": 0})])
        assert a == b

    def test_differing_extra_body_across_sources_is_distinguished(self):
        from polygateway.client import build_model_fingerprint

        a = build_model_fingerprint([_source(extra_body={"temperature": 0})])
        b = build_model_fingerprint([_source(extra_body={"temperature": 1})])
        assert a != b


class TestFactories:
    def test_from_env_assembles(self):
        client = GatewayClient.from_env("LLM", env=_ENV)
        assert isinstance(client, GatewayClient)

    def test_from_env_unknown_provider_fails_at_assembly(self):
        env = dict(_ENV)
        for key in list(env):
            if key.startswith("LLM__QWEN__"):
                env[key.replace("QWEN", "GLM")] = env.pop(key)
        with pytest.raises(ValueError, match="glm"):
            GatewayClient.from_env("LLM", env=env)

    def test_from_env_builds_redis_governance_backends(self):
        """M2: 配置取 redis 时装配出 Redis 后端(构造不连库,unit 可测)。"""
        from polygateway.backends.redis.breaker import RedisGate
        from polygateway.backends.redis.limiter import RedisLimiter
        from polygateway.client import _build_breaker, _build_limiter

        env = dict(
            _ENV,
            PGW_LIMITER_BACKEND="redis",
            PGW_BREAKER_BACKEND="redis",
            REDIS_URL="redis://:pw@10.0.0.1:6379/3",
        )
        settings = GatewaySettings.from_env("LLM", env=env)
        assert isinstance(_build_limiter(settings, list(settings.sources)), RedisLimiter)
        assert isinstance(_build_breaker(settings), RedisGate)

    def test_from_settings_respects_injection(self):
        settings = GatewaySettings.from_env("LLM", env=_ENV)
        shared = InMemoryLimiter(
            scope="shared",
            sources={s.name: s for s in settings.sources},
            global_limits=GlobalLimits(0, 0, 0),
        )
        client = GatewayClient.from_settings(settings, limiter=shared)
        assert isinstance(client, GatewayClient)


class TestSharedBackend:
    async def test_two_clients_share_global_concurrency_gate(self):
        """VT R5: 两个逻辑角色显式注入同一 limiter → 共享全局并发闸。"""
        src_a, src_b = _source("role_a_1"), _source("role_b_1")
        shared = InMemoryLimiter(
            scope="shared",
            sources={"role_a_1": src_a, "role_b_1": src_b},  # 共享后端持源并集
            global_limits=GlobalLimits(max_concurrency=1, rpm=0, tpm=0),
        )
        started = asyncio.Event()

        async def slow_handler(request):
            started.set()
            await asyncio.sleep(0.2)
            return _sse()

        client_a = _client([src_a], slow_handler, limiter=shared)
        client_b = _client([src_b], limiter=shared, quota_full="fail_fast")
        task = asyncio.ensure_future(client_a.chat([{"role": "user", "content": "x"}]))
        await started.wait()
        with pytest.raises(AllSourcesExhausted) as ei:
            await client_b.chat([{"role": "user", "content": "y"}])
        assert ei.value.reason == "quota_exhausted"  # 全局闸被 A 占满 → B 立即失败
        await task

    async def test_aclose_idempotent(self):
        client = _client()
        await client.aclose()
        await client.aclose()


class TestGatherBounded:
    async def test_order_preserved_and_concurrency_capped(self):
        peak = {"now": 0, "max": 0}

        async def work(i):
            peak["now"] += 1
            peak["max"] = max(peak["max"], peak["now"])
            await asyncio.sleep(0.01)
            peak["now"] -= 1
            return i

        results = await gather_bounded((work(i) for i in range(10)), concurrency=3)
        assert results == list(range(10))
        assert peak["max"] <= 3

    async def test_exception_propagates(self):
        async def boom():
            raise RuntimeError("x")

        async def ok():
            return 1

        with pytest.raises(RuntimeError):
            await gather_bounded([ok(), boom()], concurrency=2)

    async def test_invalid_concurrency(self):
        with pytest.raises(ValueError):
            await gather_bounded([], concurrency=0)


def _load_reference_protocol(insert_path: str, module: str):
    sys.path.insert(0, str(_REPO / insert_path))
    try:
        import importlib

        return importlib.import_module(module).LLMProvider
    finally:
        sys.path.pop(0)


class TestReferenceProtocolCompat:
    """结构兼容断言(只读 import reference;失败即公共承诺破裂)。"""

    def test_satisfies_govdoc_llm_provider(self):
        try:
            proto = _load_reference_protocol(
                "reference/GovDoc-SaaS/packages/docagent-core/src", "docagent_core.protocols"
            )
        except ImportError:
            # 兜底: 按 protocols.py:15-25 逐字复制的结构断言
            from typing import Any, Protocol, runtime_checkable

            @runtime_checkable
            class proto(Protocol):  # noqa: N801 — 复制自 GovDoc protocols.py:15-25
                async def chat(
                    self,
                    messages: list[dict[str, Any]],
                    *,
                    session_id: str | None = None,
                    parent_call_id: str | None = None,
                ): ...

        assert isinstance(_client(), proto)

    def test_satisfies_videotree_llm_provider(self):
        try:
            proto = _load_reference_protocol("reference/Video-Tree-TRM5", "core.protocols")
        except ImportError:
            from typing import Any, Protocol, runtime_checkable

            @runtime_checkable
            class proto(Protocol):  # noqa: N801 — 复制自 VT core/protocols.py:18-29
                async def chat(
                    self,
                    messages: list[dict[str, Any]],
                    *,
                    session_id: str | None = None,
                    parent_call_id: str | None = None,
                    cache_salt: str | None = None,
                ): ...

        assert isinstance(_client(), proto)