scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-08 08:08:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80
litellm/tests/integration/providers/test_gemini_long_decode_chaos.py
Mateo Wang 7090588bdf
feat(gemini)!: auto-continue gemini-4-argon long decodes with continuationToken (#46240) 
* feat(gemini): continue MAX_TOKENS decodes with continuationToken across chat, streaming, and interactions

* fix(gemini): continue long decode on CONTINUATION and cancel interactions at /cancel

* feat(gemini): add public gemini-4-argon cost-map rows for AI Studio and Vertex AI

Add gemini/gemini-4-argon and vertex_ai/gemini-4-argon alongside the existing
early-access row, carrying the same limits and capabilities. The AI Studio row
mirrors the early-access entry; the Vertex row differs only in provider and
pricing source. Prices follow once Google publishes them.

Extend the mapped metadata test to cover all three ids, that each routes to its
provider, that its backup copy matches, and that the public rows carry the same
limits as the early-access row.

* feat(gemini): raise gemini-4-argon to the 1M long decode cap and report initial_prompt_tokens across hops

* chore(prices): build on the merged gemini-4-argon rows and only raise the long decode cap

* fix(gemini): follow continuation hops in a loop instead of recursion

The unary hop collectors and the stream iterator called themselves once per continuation request, which the recursion check in CI rejects and which would hit the interpreter's call depth on a very long output. They now loop, and a test follows more hops than the interpreter allows nested calls

* fix(gemini): bill Google Search grounding from every continuation request

Google sends search grounding once, on the request that ran the search, and does not repeat it on a continuation request. The merged response kept only the last request's grounding, so usage.web_search_requests and the per-search fee dropped every earlier search. Unary merges and the stream iterator now carry the grounding of every request in the chain

* fix(gemini): continue long gemini-4-argon answers on streams and Vertex AI

A continued stream stopped after the first chunk of its second request. The stream wrapper asks the iterator for a fresh start before every chunk, and that reset it to the first request's finished stream. The iterator now reads from the newest request's stream

On Vertex AI the continuation request went to the v1 API version, which rejects continuationToken with a 400. A request that carries the token now goes to v1beta1, which accepts it, including a token that v1 issued

On a stream, initial_prompt_tokens came from a usage chunk in the middle of the stream, which reports a smaller prompt count than the last chunk. It is now the count the first request's stream ends with, the same number a unary call reports

continuation_token is no longer forwarded on the Interactions API, which answers 400 for the field. It is dropped there as it was before this branch

* chore(gemini): name the type errors the dual-mode stream iterator trips

The stream iterator holds a sync or an async line iterator in one untyped attribute. Pointing it at a continuation request's stream gave the type checker a concrete union for it, so six existing lines that only ever run in one of the two modes started to report errors

This adds a rule-named suppression with a reason on each of those lines and changes no code. Splitting the class into a sync and an async iterator would remove them and is left for a follow-up

* fix(gemini): continue only when Google ends a request with CONTINUATION, and stop streaming mid-answer tokens

A long decode now sends the next request only when Google's finish reason is CONTINUATION, so an answer cut at the caller's own max_tokens returns in one request with its continuation_token instead of drawing a rejected extra request

A request for more than one choice is never continued, since Google rejects a continuation with more than one candidate

Chat-format streams no longer carry the tokens Google sends mid-answer. Only a stream that ends cut off carries a continuation_token, on its last content chunk

The response cache key now includes continuation_token, so a resend with a token is not served the earlier segment

* fix(gemini): keep one response id across a continued answer

A continued chat stream changed its chunk id each time the next request to Google started, and a unary answer came back under the last request's id. Every chunk and the unary answer now carry the first request's id, including the chunk that ends a cut-off stream

A request whose n or max_tokens arrives as a string is sent once and never continued, where it used to fail with a 500 on the comparison. An empty continuation_token no longer changes the cache key

* test(gemini): add integration cells for long decode continuation

Covers continued answers on chat, /v1/messages and /v1/responses, streamed and not, on the Gemini API and Vertex AI against a scripted Google, plus concurrency, outage and worker restart runs

Also moves the Google stand-in cancel route to /cancel, the URL this PR now calls for Interactions cancel, and checks the cancel test posts there exactly once

* fix(vertex_ai): fetch the access token again before each continuation request

A long Gemini 4 answer on Vertex AI can run for over half an hour of follow-up requests, and every one of them reused the bearer token fetched before the first request. Once that token expired Google answered 401 and the caller got a cut-off answer with finish_reason length

Each Google request in a chain now asks for the token again, which is a cache hit until the token nears expiry, on the streamed and unary paths, sync and async. A refresh that fails counts as a failed follow-up, so the caller still gets the answer so far with its continuation token. A caller's own Authorization in extra_headers still wins on every request, and the Gemini API key path is unchanged

* fix(vertex_ai): start each continuation stream with an empty parser buffer

When a Google stream ended on an unfinished frame, its leftover text stayed in the stream parser and every frame of the next continuation stream was appended to it. None of them could be decoded, so the whole continuation was billed and dropped without an error

Each continuation stream now starts with a fresh buffer, and a test covers a stream that ends on an unfinished frame, sync and async

* fix(vertex_ai): treat a continuation stream with no readable frame as a failed follow-up

A continuation stream that answered 200 and then sent Google's error frame first, sent a body that is not JSON, or dropped before its first frame ended the caller's stream with `stop` or an error and no token, so a cut answer looked complete or was lost. It now ends the way a refused follow-up does, with the answer so far, `finish_reason` length and that request's token

Every continuation stream request now also gets the caller's per-request timeout, which #46349 added to the first one, so a follow-up that never answers ends at that timeout the same way

* fix(vertex_ai): keep every follow-up's citations and URL context in a continued answer

A non-streamed answer that took several Google requests kept citationMetadata and urlContextMetadata from the last request only, and left every later request's offsets counting from that request's own text. So a URL read on the first request was missing from the answer, citations from earlier requests were dropped, and a later request's citations and search supports pointed at the wrong part of the joined answer

The joined answer now keeps every request's citations (Vertex AI's `citations` and AI Studio's `citationSources`) and URL context, and shifts each later request's citation and search support offsets by the UTF-8 bytes of the answer text sent before it, thought parts excluded, since Google counts these offsets in bytes

* fix(vertex_ai): answer when a candidate sends null citationMetadata or urlContextMetadata

A Gemini candidate whose citationMetadata or urlContextMetadata was JSON null made the call fail on every Gemini model, non-streamed with a 422 or 500 and streamed with an error event, because the per-request metadata helper handed the null back as is and the caller tried to extend a list with it. A continued answer also failed when a follow-up's citation or search support came with a null startIndex

A null field is now kept as one entry on a single request, as it was before this branch, a follow-up's null URL context adds nothing to a continued answer, and a null startIndex counts from the start of that request's text. Tests cover missing, null and empty fields on one request and on a continued answer, streamed and not

* fix(vertex_ai): pass through spans with no end index and point a continued stream's citations into the whole answer

A later request's citation or search support with no endIndex made the whole non-streamed answer fail with a 500, so a span Google sends without an end index now comes back as Google sent it, the way one request for the whole answer returns it. A span with an end index but no start index still counts from 0

A continued stream handed callers each later request's citation and search support offsets counted from that request's own text. The stream now adds the answer bytes of every earlier request, so the offsets point into the whole answer like the non-streamed path and one request for the whole answer

Adds the audit's wire tests for both, which fail 24 cells before this change

* test(vertex_ai): cover a later request whose cited span starts at its first byte with no start index

Google leaves out a zero startIndex, so the second request of a continued answer can cite its own opening with only an endIndex. These wire cells check that span lands where one request for the whole answer puts it, on Gemini and Vertex AI, chat and Responses, unary and streamed

* fix(vertex_ai): log a continued answer's raw response as one object

A non-streamed answer joined from several Google requests logged its raw response as a JSON list of every request's body. The OpenTelemetry logger reads that raw response as one Gemini body to set its `llm.gemini.*` attributes, so on a continued answer it hit a validation error and set none of them

It now logs the merged body, which carries only the last request's token, and a single request still logs Google's body unchanged. The new unit test fails with the change reverted

Also adds integration cells for a cut chat stream's token chunk, an answer joined from 41 requests, a caller that gave up on a non-streamed call, counts and caps that are not numbers or are sent twice, `n` sent as true, a partial answer cached after a refused continuation, the token in another provider's cache key, and the sync SDK path with its Vertex token refetch

* test(integration): check a continued answer logs one merged raw response

A non-streamed Gemini answer that took several Google requests now hands logging one merged response, so the OpenTelemetry span carries the llm.gemini attributes. These 12 cells fail on the merge base and on the commit before that fix, and pass on the branch

* test(vertex_ai): cover a follow-up tool call and its thought signature in long decode

* test(gemini): pin that a one-request answer logs Google's body unchanged
2026-10-15 22:54:23 +00:00

470 lines
18 KiB
Python
Raw Blame History

from __future__ import annotations
import asyncio
import dataclasses
import re
import threading
import uuid
from collections.abc import Callable, Generator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final
from urllib.parse import urlsplit
import httpx
import psutil
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, eventually, string_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.responses_vendor import same_response
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.providers._gemini_long_decode_support import (
API_KEY,
BACKEND,
CHECKPOINT,
ENDPOINTS,
INPUT_COST,
OUTPUT_COST,
PLAIN,
PROVIDERS,
SPEND_COLUMNS,
Endpoint,
Hop,
Seen,
call_type,
carries_token,
complete,
cost_of,
cut_off,
hop_reply,
number,
path_of,
payloads,
post,
register,
request_body,
seen,
single,
streams,
summed,
)
from pydantic import JsonValue
Continuation = Callable[[Request, str], Reply]
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_HOP_TOKEN: Final = "hop-token-"
_PROBE_ID: Final = "probe-"
_MODES: Final = len(ENDPOINTS) * 2
@dataclass(frozen=True, slots=True)
class Call:
endpoint: Endpoint
stream: bool
marker: str
model: str
@dataclass(frozen=True, slots=True)
class Served:
call: Call
status: int
text: str
@property
def observed(self) -> Seen:
assert self.status == 200, self.text
return seen(self.call.endpoint, self.call.stream, payloads(self.text, self.call.stream), self.text)
def _opening(marker: str) -> Hop:
return Hop((f"opening marker-{marker} ", "and "), "CONTINUATION", f"{_HOP_TOKEN}{marker}", 20, 7)
def _closing(marker: str) -> Hop:
return Hop((f"closing marker-{marker}", "."), "STOP", None, 27, 5)
def _joined(marker: str) -> str:
return _opening(marker).text + _closing(marker).text
_OPENING_USAGE: Final = _opening("").usage
_CHAIN_USAGE: Final = summed(_opening(""), _closing(""))
def _marker_asked(request: Request) -> str | None:
found: Final = _MARKER.search(request.body.decode())
return found.group(1) if found is not None else None
def _token_sent(request: Request) -> str | None:
token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken")
return token if isinstance(token, str) else None
def _continues(request: Request, token: str) -> Reply:
return hop_reply(f"second-{uuid.uuid4().hex}", _closing(token.removeprefix(_HOP_TOKEN)), streams(request))
def _continues_slowly(request: Request, token: str) -> Reply:
return dataclasses.replace(_continues(request, token), pause_between_chunks=0.3)
def _held(release: threading.Event, held: SimpleQueue[str], then: Continuation = _continues) -> Continuation:
def hold(request: Request, token: str) -> Reply:
held.put(token)
assert release.wait(timeout=90), "The continuation requests were never released"
return then(request, token)
return hold
def _down_for(doomed: frozenset[str]) -> Continuation:
def answer(request: Request, token: str) -> Reply:
if token.removeprefix(_HOP_TOKEN) in doomed:
return Reply(drop_connection=True)
return _continues(request, token)
return answer
def _google(continuation: Continuation = _continues) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
marker: Final = _marker_asked(request)
token: Final = _token_sent(request)
if marker is None:
return hop_reply(f"{_PROBE_ID}{uuid.uuid4().hex}", PLAIN, streams(request))
if token is None:
return hop_reply(f"first-{marker}", _opening(marker), streams(request))
return continuation(request, token)
return respond
@contextmanager
def _released_on_exit(release: threading.Event) -> Generator[threading.Event]:
try:
yield release
finally:
release.set()
def _calls(count: int, models: Sequence[str], stream: Callable[[int], bool]) -> tuple[Call, ...]:
return tuple(
Call(
endpoint=ENDPOINTS[index % len(ENDPOINTS)],
stream=stream(index),
marker=uuid.uuid4().hex,
model=models[(index // _MODES) % len(models)],
)
for index in range(count)
)
def _alternating(index: int) -> bool:
return index % 2 == 0
async def _send(client: httpx.AsyncClient, key: str, call: Call) -> Served:
question: Final = f"Write the long answer for marker-{call.marker}"
async with client.stream(
"POST",
path_of(call.endpoint),
json=request_body(call.endpoint, call.model, call.stream, prompt=question),
headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
) as response:
raw: Final = await response.aread()
return Served(call=call, status=response.status_code, text=raw.decode())
async def _burst(
gateway: Gateway, calls: Sequence[Call], *, tolerate_transport_errors: bool = False
) -> tuple[Served, ...]:
async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=90, trust_env=False) as client:
results: Final = await asyncio.gather(
*(_send(client, gateway.key, call) for call in calls), return_exceptions=tolerate_transport_errors
)
for result in results:
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
return tuple(result for result in results if isinstance(result, Served))
async def _all_mid_chain(held: SimpleQueue[str], burst: asyncio.Task[tuple[Served, ...]], count: int) -> None:
await asyncio.to_thread(eventually, lambda: held.qsize() == count or burst.done(), bool, 60)
assert held.qsize() == count, f"{held.qsize()} of {count} answers reached their continuation request"
def _assert_joined(served: Served) -> None:
observed: Final = served.observed
whole: Final = _joined(served.call.marker)
assert (observed.answer, observed.streamed) == (whole, whole), served.text
assert observed.finishes == (complete(served.call.endpoint),), served.text
assert (observed.usage, observed.tokens, observed.errors) == (_CHAIN_USAGE, (), ()), served.text
assert CHECKPOINT not in served.text and _HOP_TOKEN not in served.text, served.text
assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text
def _assert_cut_off_with_its_own_token(served: Served) -> None:
observed: Final = served.observed
call: Final = served.call
so_far: Final = _opening(call.marker).text
own: Final = (f"{_HOP_TOKEN}{call.marker}",) if carries_token(call.endpoint, call.stream) else ()
assert (observed.answer, observed.streamed) == (so_far, so_far), served.text
assert observed.finishes == (cut_off(call.endpoint),), served.text
assert (observed.usage, observed.tokens, observed.errors) == (_OPENING_USAGE, own, ()), served.text
assert CHECKPOINT not in served.text, served.text
assert set(_MARKER.findall(served.text)) == {call.marker}, served.text
def _assert_each_answer_sent_its_own_token_once(received: Sequence[Request], calls: Sequence[Call]) -> None:
sent: Final = sorted(
(marker, _token_sent(request) or "") for request in received if (marker := _marker_asked(request)) is not None
)
assert sent == sorted(
(*((call.marker, "") for call in calls), *((call.marker, f"{_HOP_TOKEN}{call.marker}") for call in calls))
), sent
def _spend_rows(model: str) -> Sequence[Mapping[str, JsonValue]]:
return read_rows(
f"{SPEND_COLUMNS} WHERE model_group=%s AND request_id NOT LIKE '{_PROBE_ID}%%' ORDER BY request_id", (model,)
)
def _billed_as(row: Mapping[str, JsonValue], served: Served, usage: tuple[int, int]) -> bool:
identities: Final = (single(served.observed.ids), f"first-{served.call.marker}")
return (
any(same_response(string_value(row["request_id"]), identity) for identity in identities)
and (row["call_type"], row["status"], row["prompt_tokens"], row["completion_tokens"])
== (call_type(served.call.endpoint), "success", *usage)
and number(row["spend"]) == pytest.approx(cost_of(usage))
)
def _billings(rows: Sequence[Mapping[str, JsonValue]], served: Served, usage: tuple[int, int]) -> int:
return sum(1 for row in rows if _billed_as(row, served, usage))
def _assert_landed_once(model: str, joined: Sequence[Served], cut_off_early: Sequence[Served] = ()) -> None:
mine: Final = tuple(
(served, usage)
for served, usage in (
*((served, _CHAIN_USAGE) for served in joined),
*((served, _OPENING_USAGE) for served in cut_off_early),
)
if served.call.model == model
)
rows: Final = eventually(lambda: _spend_rows(model), lambda found: len(found) >= len(mine), seconds=90)
assert len(rows) == len(mine), (len(rows), len(mine), rows)
unbilled: Final = tuple(served.call for served, usage in mine if _billings(rows, served, usage) != 1)
assert unbilled == (), (unbilled, rows)
@pytest.mark.timeout(400)
async def test_mixed_burst_of_continued_answers_keeps_every_answer_on_its_own_token_and_bills_it_once(
gateway: Gateway,
) -> None:
release: Final = threading.Event()
held: Final[SimpleQueue[str]] = SimpleQueue()
with (
wire_server(_google(_held(release, held))) as wire,
_released_on_exit(release),
gateway.scenario() as scenario,
):
models: Final = tuple(register(gateway, scenario, provider, wire) for provider in PROVIDERS)
calls: Final = _calls(30, models, _alternating)
burst: Final = asyncio.create_task(_burst(gateway, calls))
await _all_mid_chain(held, burst, 30)
release.set()
served: Final = await burst
assert len(served) == 30
for item in served:
_assert_joined(item)
_assert_each_answer_sent_its_own_token_once(wire.drain(), calls)
for model in models:
_assert_landed_once(model, served)
@pytest.mark.timeout(400)
async def test_slow_continuation_streams_are_joined_once_each_under_concurrency(gateway: Gateway) -> None:
with wire_server(_google(_continues_slowly)) as wire, gateway.scenario() as scenario:
model: Final = register(gateway, scenario, "gemini", wire)
calls: Final = _calls(20, (model,), lambda _: True)
served: Final = await _burst(gateway, calls)
assert len(served) == 20
for item in served:
_assert_joined(item)
_assert_each_answer_sent_its_own_token_once(wire.drain(), calls)
_assert_landed_once(model, served)
@pytest.mark.timeout(400)
async def test_google_outage_on_half_the_continuation_requests_mid_burst_cuts_those_off_and_recovers(
gateway: Gateway,
) -> None:
release: Final = threading.Event()
held: Final[SimpleQueue[str]] = SimpleQueue()
with gateway.scenario() as scenario:
calls: Final = _calls(30, ("unregistered",), _alternating)
doomed: Final = frozenset(call.marker for call in calls[::2])
with wire_server(_google(_held(release, held, _down_for(doomed)))) as wire, _released_on_exit(release):
model: Final = register(gateway, scenario, "gemini", wire)
routed: Final = tuple(dataclasses.replace(call, model=model) for call in calls)
burst: Final = asyncio.create_task(_burst(gateway, routed))
await _all_mid_chain(held, burst, 30)
release.set()
served: Final = await burst
assert len(served) == 30
lost: Final = tuple(item for item in served if item.call.marker in doomed)
kept: Final = tuple(item for item in served if item.call.marker not in doomed)
for item in lost:
_assert_cut_off_with_its_own_token(item)
for item in kept:
_assert_joined(item)
after: Final = _calls(_MODES, (model,), _alternating)
recovered: Final = await _burst(gateway, after)
assert len(recovered) == _MODES
for item in recovered:
_assert_joined(item)
_assert_each_answer_sent_its_own_token_once(wire.drain(), (*routed, *after))
_assert_landed_once(model, (*kept, *recovered), lost)
def _chaos_config(wire: Wire, directory: Path, model_name: str) -> Path:
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = {
**base,
"model_list": [
{
"model_name": model_name,
"litellm_params": {
"model": f"gemini/{BACKEND}",
"api_base": wire.url,
"api_key": API_KEY,
"input_cost_per_token": INPUT_COST,
"output_cost_per_token": OUTPUT_COST,
},
}
],
}
path: Final = directory / "gemini-long-decode-chaos.yaml"
path.write_text(yaml.safe_dump(config))
return path
def _answers_probes(gateway: Gateway, model: str) -> bool:
try:
return post(gateway, "chat", request_body("chat", model, False)).status_code == 200
except httpx.TransportError:
return False
def _ready(gateway: Gateway, model: str) -> None:
eventually(
lambda: tuple(_answers_probes(gateway, model) for _ in range(2)), all, seconds=graceful_stop_seconds() + 60
)
def _open_upstream_connections(pid: int, upstream: str) -> int:
port: Final = urlsplit(upstream).port
return sum(
1
for connection in psutil.Process(pid).net_connections(kind="tcp")
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
)
def _worker_pids(log: Path) -> tuple[int, ...]:
return eventually(
lambda: tuple(int(started.group(1)) for started in _STARTED_WORKER.finditer(log.read_text())),
lambda pids: len(pids) == 2,
seconds=graceful_stop_seconds() + 60,
)
@pytest.mark.timeout(2 * graceful_stop_seconds() + 240)
async def test_worker_sigkill_mid_answer_leaves_the_sibling_joining_its_own_answers(
gateway: Gateway, tmp_path: Path
) -> None:
release: Final = threading.Event()
held: Final[SimpleQueue[str]] = SimpleQueue()
alias: Final = f"gemini-long-decode-chaos-{uuid.uuid4().hex[:12]}"
calls: Final = _calls(20, (alias,), _alternating)
with wire_server(_google(_held(release, held))) as wire, _released_on_exit(release):
config: Final = _chaos_config(wire, tmp_path, alias)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
candidate: Final = owned.gateway
workers: Final = _worker_pids(owned.log)
_ready(candidate, alias)
wire.drain()
burst: Final = asyncio.create_task(_burst(candidate, calls, tolerate_transport_errors=True))
await _all_mid_chain(held, burst, 20)
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
assert sum(held_by.values()) == 20, held_by
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
with owned.replacing_worker(victim_pid) as kill:
kill.freeze()
await asyncio.to_thread(kill)
release.set()
served: Final = await burst
assert held_by[survivor_pid] >= 10, held_by
assert len(served) == held_by[survivor_pid], (held_by, len(served))
for item in served:
_assert_joined(item)
follow_up: Final = _calls(1, (alias,), _alternating)
(answered,) = await _burst(candidate, follow_up)
_assert_joined(answered)
_assert_each_answer_sent_its_own_token_once(wire.drain(), (*calls, *follow_up))
_assert_landed_once(alias, (*served, answered))
eventually(
lambda: owned.log.read_text().count("Application startup complete."),
lambda started: started == 3,
seconds=graceful_stop_seconds() + 60,
)
@pytest.mark.timeout(4 * graceful_stop_seconds() + 240)
async def test_proxy_sigterm_mid_answer_finishes_every_answer_and_the_rebooted_proxy_bills_each_once(
gateway: Gateway, tmp_path: Path
) -> None:
release: Final = threading.Event()
held: Final[SimpleQueue[str]] = SimpleQueue()
alias: Final = f"gemini-long-decode-chaos-{uuid.uuid4().hex[:12]}"
calls: Final = _calls(20, (alias,), _alternating)
with wire_server(_google(_held(release, held))) as wire, _released_on_exit(release):
config: Final = _chaos_config(wire, tmp_path, alias)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
_worker_pids(owned.log)
_ready(owned.gateway, alias)
wire.drain()
burst: Final = asyncio.create_task(_burst(owned.gateway, calls))
await _all_mid_chain(held, burst, 20)
owned.begin_shutdown()
release.set()
served: Final = await burst
assert len(served) == 20
for item in served:
_assert_joined(item)
await asyncio.to_thread(
eventually, owned.process.poll, lambda code: code is not None, graceful_stop_seconds() + 60
)
_assert_each_answer_sent_its_own_token_once(wire.drain(), calls)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as rebooted:
_worker_pids(rebooted.log)
_ready(rebooted.gateway, alias)
wire.drain()
follow_up: Final = _calls(_MODES, (alias,), _alternating)
recovered: Final = await _burst(rebooted.gateway, follow_up)
assert len(recovered) == _MODES
for item in recovered:
_assert_joined(item)
_assert_each_answer_sent_its_own_token_once(wire.drain(), follow_up)
_assert_landed_once(alias, (*served, *recovered))
Reference in a new issue View git blame Copy permalink