SilverElixir commited on
Commit
fb2a2f0
·
1 Parent(s): 80c432a

Fix Gemini streaming await and correct SDK contract mocks

Browse files
Files changed (2) hide show
  1. lumen_streaming.py +4 -7
  2. tests/test_bot_streaming.py +12 -14
lumen_streaming.py CHANGED
@@ -94,14 +94,11 @@ def _stream_wait(bot: Any, deadline: float | None) -> float:
94
  async def _gemini_stream_pieces(model_id: str, call_contents: list, gconfig, *, deadline: float | None = None):
95
  """Куски от Gemini (тонкая обёртка над generate_content_stream).
96
 
97
- Рукопожатия как отдельного шага НЕТ намеренно: generate_content_stream
98
- возвращает генератор, не делая сети (проверено по исходнику google-genai
99
- 2.24.0: тело — `return stream_generator()`), поэтому wait_for вокруг вызова
100
- ограничивал только мгновенное создание объекта и никогда не срабатывал —
101
- витрина таймаута без защиты (враждебное ревью 27.09.2026). Реальную сетевую
102
- жду ловлю ниже: первый кусок идёт через __anext__ под тем же _stream_wait."""
103
  import bot
104
- stream = bot.client.aio.models.generate_content_stream(model=model_id, contents=call_contents, config=gconfig)
105
  stream_iter = stream.__aiter__()
106
  try:
107
  while True:
 
94
  async def _gemini_stream_pieces(model_id: str, call_contents: list, gconfig, *, deadline: float | None = None):
95
  """Куски от Gemini (тонкая обёртка над generate_content_stream).
96
 
97
+ В google-genai 2.27.0 это async def: вызов надо ждать (пример из SDK:
98
+ `async for chunk in await client.aio.models.generate_content_stream(...)`).
99
+ Реальную сетевую жду ловлю ниже: первый кусок идёт через __anext__ под тем же _stream_wait."""
 
 
 
100
  import bot
101
+ stream = await bot.client.aio.models.generate_content_stream(model=model_id, contents=call_contents, config=gconfig)
102
  stream_iter = stream.__aiter__()
103
  try:
104
  while True:
tests/test_bot_streaming.py CHANGED
@@ -35,15 +35,10 @@ def test_stream_wait_caps_chunk_timeout_by_remaining_budget():
35
 
36
 
37
  def _fake_gemini_stream(pieces=None, *, raises=None, hang_sec=0.0):
38
- """Мок generate_content_stream по НАСТОЯЩЕМУ контракту google-genai 2.24.0:
39
- обычная функция, возвращающая асинхронный генератор (тело SDK — `return
40
- stream_generator()`), а НЕ корутина.
41
-
42
- Раньше фейки здесь были `async def ... return gen()`, из-за чего тесты
43
- подпирали фиктивное рукопожатие wait_for вокруг вызова: реальный SDK сети
44
- в корутине не делает, значит тот таймаут не срабатывал никогда
45
- (враждебное ревью 27.09.2026)."""
46
- def _stream(*, model, contents, config=None):
47
  async def gen():
48
  if hang_sec:
49
  await asyncio.sleep(hang_sec)
@@ -76,19 +71,22 @@ def test_gemini_stream_timeout_covers_a_hanging_first_chunk():
76
  bot.STREAM_CHUNK_TIMEOUT_SEC = original_cap
77
 
78
 
79
- def test_gemini_stream_call_is_not_awaited():
80
- # Контракт SDK: generate_content_stream — обычная функция. Если бы обёртка снова
81
- # стала await-ить её, тест падал бы с TypeError, а не проходил вхолостую.
 
 
 
82
  calls = []
83
 
84
- def sync_only(*, model, contents, config=None):
85
  calls.append(model)
86
  async def gen():
87
  yield SimpleNamespace(text="ok")
88
  return gen()
89
 
90
  fake_client = MagicMock()
91
- fake_client.aio.models.generate_content_stream = sync_only
92
  original_client = bot.client
93
  bot.client = fake_client
94
  try:
 
35
 
36
 
37
  def _fake_gemini_stream(pieces=None, *, raises=None, hang_sec=0.0):
38
+ """Мок generate_content_stream по контракту google-genai 2.27.0:
39
+ async def, вызов надо ждать (пример из SDK:
40
+ `async for chunk in await client.aio.models.generate_content_stream(...)`)."""
41
+ async def _stream(*, model, contents, config=None):
 
 
 
 
 
42
  async def gen():
43
  if hang_sec:
44
  await asyncio.sleep(hang_sec)
 
71
  bot.STREAM_CHUNK_TIMEOUT_SEC = original_cap
72
 
73
 
74
+ def test_gemini_stream_call_is_awaited():
75
+ # Контракт SDK 2.27.0: generate_content_stream это async def, обёртка обязана
76
+ # ждать вызов. Синхронный мок (старое поведение) даёт TypeError при await.
77
+ import inspect
78
+ from google.genai.models import AsyncModels
79
+ assert inspect.iscoroutinefunction(AsyncModels.generate_content_stream)
80
  calls = []
81
 
82
+ async def async_only(*, model, contents, config=None):
83
  calls.append(model)
84
  async def gen():
85
  yield SimpleNamespace(text="ok")
86
  return gen()
87
 
88
  fake_client = MagicMock()
89
+ fake_client.aio.models.generate_content_stream = async_only
90
  original_client = bot.client
91
  bot.client = fake_client
92
  try: