"""三层流式看门狗 (stream_with_liveness_timeouts) 单元测试。 验证 TTFT / inter_token / total 三层超时保护行为, 以及正常透传、空流等边界场景。 """ from __future__ import annotations import asyncio from typing import TYPE_CHECKING, TypeVar if TYPE_CHECKING: from collections.abc import AsyncIterator import pytest from adapters.streaming import StreamLivenessTimeout, stream_with_liveness_timeouts _T = TypeVar("_T") async def _items_with_delays( items: list[_T], delays: list[float], ) -> AsyncIterator[_T]: """辅助异步生成器:按 delays 指定的延迟逐项产出 items。 参数: items: 待产出的元素列表。 delays: 每个元素产出前的延迟秒数,长度须与 items 一致。 """ for item, delay in zip(items, delays, strict=True): await asyncio.sleep(delay) yield item @pytest.mark.asyncio async def test_all_items_yielded_within_budget() -> None: """正常场景:所有 item 在超时预算内产出,全部透传。""" items = ["a", "b", "c", "d"] delays = [0.01, 0.01, 0.01, 0.01] result: list[str] = [] async for chunk in stream_with_liveness_timeouts( _items_with_delays(items, delays), ttft_s=1.0, inter_token_s=1.0, total_s=5.0, ): result.append(chunk) assert result == items @pytest.mark.asyncio async def test_ttft_timeout_fires() -> None: """首 token 超时:第一个 item 延迟超过 ttft_s,触发 kind='ttft'。""" items = ["slow_first"] delays = [0.5] with pytest.raises(StreamLivenessTimeout) as exc_info: async for _ in stream_with_liveness_timeouts( _items_with_delays(items, delays), ttft_s=0.05, inter_token_s=5.0, total_s=10.0, ): pass # pragma: no cover — 不应到达 err = exc_info.value assert err.kind == "ttft" assert err.first_token_seen is False @pytest.mark.asyncio async def test_inter_token_timeout_fires() -> None: """token 间隔超时:首项正常,第二项延迟超过 inter_token_s。""" items = ["fast", "very_slow"] delays = [0.01, 0.5] with pytest.raises(StreamLivenessTimeout) as exc_info: async for _ in stream_with_liveness_timeouts( _items_with_delays(items, delays), ttft_s=5.0, inter_token_s=0.05, total_s=10.0, ): pass err = exc_info.value assert err.kind == "inter_token" assert err.first_token_seen is True @pytest.mark.asyncio async def test_total_timeout_fires() -> None: """总时长超时:每个 item 都在单层预算内,但累计超过 total_s。""" items = ["a", "b", "c", "d", "e", "f"] # 每项 0.05s,累计 0.30s > total_s=0.15s delays = [0.05, 0.05, 0.05, 0.05, 0.05, 0.05] with pytest.raises(StreamLivenessTimeout) as exc_info: async for _ in stream_with_liveness_timeouts( _items_with_delays(items, delays), ttft_s=5.0, inter_token_s=5.0, total_s=0.15, ): pass err = exc_info.value assert err.kind == "total" @pytest.mark.asyncio async def test_thinking_tokens_refresh_watchdog() -> None: """(False, 'think') 元组正常透传,不影响看门狗计时。""" items = [ (False, "think_1"), (False, "think_2"), (True, "real_token"), ] delays = [0.01, 0.01, 0.01] result: list[tuple[bool, str]] = [] async for chunk in stream_with_liveness_timeouts( _items_with_delays(items, delays), ttft_s=1.0, inter_token_s=1.0, total_s=5.0, ): result.append(chunk) assert result == items @pytest.mark.asyncio async def test_empty_stream_no_error() -> None: """空流:source 立即耗尽,不抛异常,正常结束。""" items: list[str] = [] delays: list[float] = [] result: list[str] = [] async for chunk in stream_with_liveness_timeouts( _items_with_delays(items, delays), ttft_s=1.0, inter_token_s=1.0, total_s=5.0, ): result.append(chunk) # pragma: no cover assert result == []