test(integration): cover outbound TLS cipher list at the handler and proxy edges

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-24 08:06:13 +00:00
parent dea88965fb
commit 29945f1b66
3 changed files with 696 additions and 0 deletions

View file

@ -0,0 +1,194 @@
from __future__ import annotations
import datetime
import ipaddress
import json
import ssl
import threading
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final, cast
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID
from pydantic import JsonValue, TypeAdapter
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def write_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
now: Final = datetime.datetime.now(datetime.timezone.utc)
name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")])
cert: Final = (
x509.CertificateBuilder()
.subject_name(name)
.issuer_name(name)
.public_key(key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now - datetime.timedelta(days=1))
.not_valid_after(now + datetime.timedelta(days=7))
.add_extension(
x509.SubjectAlternativeName([x509.DNSName("localhost"), x509.IPAddress(ipaddress.ip_address("127.0.0.1"))]),
critical=False,
)
.sign(key, hashes.SHA256())
)
cert_file: Final = cert_dir / "cert.pem"
key_file: Final = cert_dir / "key.pem"
cert_file.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
key_file.write_bytes(
key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.TraditionalOpenSSL,
serialization.NoEncryption(),
)
)
return cert_file, key_file
@dataclass(frozen=True, slots=True)
class TlsPeerRequest:
path: str
user: str | None
cipher: str
def _completion_body(identifier: str, model: str) -> dict[str, JsonValue]:
return {
"id": identifier,
"object": "chat.completion",
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "tls ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
def _chunk_body(
identifier: str, model: str, content: str, role: str | None, finish_reason: str | None
) -> dict[str, JsonValue]:
delta: dict[str, JsonValue] = {"content": content}
if role is not None:
delta["role"] = role
return {
"id": identifier,
"object": "chat.completion.chunk",
"model": model,
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
}
class _CipherPeer(ThreadingHTTPServer):
daemon_threads = True
allow_reuse_address = True
def __init__(self, address: tuple[str, int], owner: TlsPeer) -> None:
super().__init__(address, _CipherPeerHandler)
self.owner = owner
def _negotiated_cipher(handler: BaseHTTPRequestHandler) -> str:
info: Final = cast(ssl.SSLSocket, handler.connection).cipher()
return info[0] if info else ""
class _CipherPeerHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
peer: Final = cast(_CipherPeer, self.server).owner
body: Final = JSON_OBJECT.validate_json(self.rfile.read(int(self.headers["Content-Length"])))
user_value: Final = body.get("user")
user: Final = user_value if isinstance(user_value, str) else None
peer.record(TlsPeerRequest(path=self.path, user=user, cipher=_negotiated_cipher(self)))
identifier: Final = f"chatcmpl-{user or ''}"
model_value: Final = body.get("model")
model: Final = model_value if isinstance(model_value, str) else "gpt-4o-mini"
if body.get("stream") is True:
chunks: Final = (
f"data: {json.dumps(_chunk_body(identifier, model, 'tls', 'assistant', None))}\n\n"
+ f"data: {json.dumps(_chunk_body(identifier, model, ' ok', None, 'stop'))}\n\n"
+ "data: [DONE]\n\n"
).encode()
self.send_response(200)
self.send_header("Content-Type", "text/event-stream")
self.send_header("Content-Length", str(len(chunks)))
self.end_headers()
self.wfile.write(chunks)
return
payload: Final = json.dumps(_completion_body(identifier, model)).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, format: str, *args: object) -> None:
return
class TlsPeer:
def __init__(self, cipher: str, cert_file: Path, key_file: Path) -> None:
self.cipher = cipher
self.cert_file = cert_file
self.key_file = key_file
self.port = 0
self._server: _CipherPeer | None = None
self._thread: threading.Thread | None = None
self._lock = threading.Lock()
self._requests: list[TlsPeerRequest] = []
@property
def url(self) -> str:
return f"https://127.0.0.1:{self.port}"
def record(self, request: TlsPeerRequest) -> None:
with self._lock:
self._requests.append(request)
def received(self) -> tuple[TlsPeerRequest, ...]:
with self._lock:
return tuple(self._requests)
def start(self) -> None:
server: Final = _CipherPeer(("127.0.0.1", self.port), self)
context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
context.load_cert_chain(str(self.cert_file), str(self.key_file))
context.maximum_version = ssl.TLSVersion.TLSv1_2
context.set_ciphers(self.cipher)
server.socket = context.wrap_socket(server.socket, server_side=True)
self.port = server.server_address[1]
self._server = server
self._thread = threading.Thread(target=server.serve_forever, daemon=True)
self._thread.start()
def stop(self) -> None:
if self._server is None:
return
self._server.shutdown()
self._server.server_close()
if self._thread is not None:
self._thread.join(timeout=10)
self._server = None
self._thread = None
@contextmanager
def tls_peer(cipher: str, cert_file: Path, key_file: Path) -> Generator[TlsPeer, None, None]:
peer: Final = TlsPeer(cipher, cert_file, key_file)
peer.start()
try:
yield peer
finally:
peer.stop()

View file

@ -0,0 +1,188 @@
from __future__ import annotations
import asyncio
import uuid
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import pytest
import yaml
from openai import APIStatusError, AsyncOpenAI, OpenAI
from pydantic import JsonValue, TypeAdapter
from tests.integration._support.client import Gateway, eventually, gateway_from_environment
from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy
from tests.integration._support.tls import TlsPeer, tls_peer, write_self_signed_cert
CHACHA20: Final = "ECDHE-RSA-CHACHA20-POLY1305"
NONPFS_GCM: Final = "AES256-GCM-SHA384"
PFS_GCM: Final = "ECDHE-RSA-AES256-GCM-SHA384"
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
@dataclass(frozen=True, slots=True)
class Peers:
gcm: TlsPeer
chacha: TlsPeer
nonpfs: TlsPeer
@pytest.fixture(scope="module")
def tls_cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]:
return write_self_signed_cert(tmp_path_factory.mktemp("tlscerts"))
@pytest.fixture(scope="module")
def tls_peers(tls_cert: tuple[Path, Path]) -> Iterator[Peers]:
cert_file, key_file = tls_cert
with (
tls_peer(PFS_GCM, cert_file, key_file) as gcm,
tls_peer(CHACHA20, cert_file, key_file) as chacha,
tls_peer(NONPFS_GCM, cert_file, key_file) as nonpfs,
):
yield Peers(gcm=gcm, chacha=chacha, nonpfs=nonpfs)
@pytest.fixture(scope="module")
def cipher_proxy(
tls_cert: tuple[Path, Path], tls_peers: Peers, tmp_path_factory: pytest.TempPathFactory
) -> Iterator[Gateway]:
cert_file, _ = tls_cert
directory: Final = tmp_path_factory.mktemp("tlsproxy")
base_config: Final = JSON_OBJECT.validate_python(
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
)
config: Final = {
**base_config,
"model_list": [
{
"model_name": "tls-peer",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{tls_peers.gcm.url}/v1",
"api_key": "sk-peer",
},
},
{
"model_name": "tls-chacha-peer",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{tls_peers.chacha.url}/v1",
"api_key": "sk-peer",
},
},
{
"model_name": "tls-nonpfs-peer",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{tls_peers.nonpfs.url}/v1",
"api_key": "sk-peer",
},
},
],
}
config_path: Final = directory / "tls_cipher_config.yaml"
config_path.write_text(yaml.safe_dump(config))
with gateway_from_environment() as gateway:
with owned_proxy(gateway, directory, {"SSL_CERT_FILE": str(cert_file)}, config=config_path) as candidate:
yield candidate
def _sync_completion(proxy: Gateway, model: str, user: str, *, stream: bool = False) -> str:
with OpenAI(api_key=proxy.key, base_url=f"{proxy.client.base_url}/v1", max_retries=0, timeout=60) as client:
if stream:
chunks: Final = list(
client.chat.completions.create(
model=model, messages=[{"role": "user", "content": "tls"}], user=user, stream=True
)
)
assert chunks, "stream produced no chunks"
assert all(chunk.id == f"chatcmpl-{user}" for chunk in chunks), chunks[0].id
return str(chunks[0].id)
response: Final = client.chat.completions.create(
model=model, messages=[{"role": "user", "content": "tls"}], user=user
)
return str(response.id)
async def _async_completion(proxy: Gateway, model: str, user: str, *, stream: bool = False) -> str:
async with AsyncOpenAI(
api_key=proxy.key, base_url=f"{proxy.client.base_url}/v1", max_retries=0, timeout=60
) as client:
if stream:
chunk_ids: Final = [
chunk.id
async for chunk in await client.chat.completions.create(
model=model, messages=[{"role": "user", "content": "tls"}], user=user, stream=True
)
]
assert chunk_ids, "stream produced no chunks"
assert all(chunk_id == f"chatcmpl-{user}" for chunk_id in chunk_ids), chunk_ids[0]
return str(chunk_ids[0])
response: Final = await client.chat.completions.create(
model=model, messages=[{"role": "user", "content": "tls"}], user=user
)
return str(response.id)
def test_proxy_round_trips_over_pfs_gcm_peer_and_writes_spend(cipher_proxy: Gateway, tls_peers: Peers) -> None:
marker: Final = uuid.uuid4().hex[:8]
sync_id: Final = _sync_completion(cipher_proxy, "tls-peer", f"user13-sync-{marker}")
stream_id: Final = _sync_completion(cipher_proxy, "tls-peer", f"user13-stream-{marker}", stream=True)
async_id: Final = asyncio.run(_async_completion(cipher_proxy, "tls-peer", f"user13-async-{marker}"))
async_stream_id: Final = asyncio.run(
_async_completion(cipher_proxy, "tls-peer", f"user13-astream-{marker}", stream=True)
)
users: Final = {f"user13-{leg}-{marker}" for leg in ("sync", "stream", "async", "astream")}
recorded: Final = tls_peers.gcm.received()
assert {record.user for record in recorded} == users, recorded
assert all(record.cipher == PFS_GCM for record in recorded), recorded
spend_rows: Final = tuple(
eventually(
lambda response_id=response_id: read_rows(
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)
),
lambda values: len(values) == 1,
seconds=70,
)
for response_id in (sync_id, stream_id, async_id, async_stream_id)
)
assert [rows[0]["request_id"] for rows in spend_rows] == [sync_id, stream_id, async_id, async_stream_id]
def test_proxy_reports_ssl_error_for_chacha20_only_peer(cipher_proxy: Gateway, tls_peers: Peers) -> None:
marker: Final = uuid.uuid4().hex[:8]
before: Final = len(tls_peers.chacha.received())
with pytest.raises(APIStatusError) as raised:
_sync_completion(cipher_proxy, "tls-chacha-peer", f"user14-{marker}")
assert raised.value.status_code >= 400
assert "ssl" in str(raised.value).lower() or "connection error" in str(raised.value).lower(), str(raised.value)
assert len(tls_peers.chacha.received()) == before
readiness: Final = cipher_proxy.client.get("/health/readiness")
assert readiness.status_code == 200, readiness.text
recovered: Final = _sync_completion(cipher_proxy, "tls-peer", f"user14-ok-{marker}")
assert recovered == f"chatcmpl-user14-ok-{marker}"
def test_proxy_reports_ssl_error_for_non_pfs_peer(cipher_proxy: Gateway, tls_peers: Peers) -> None:
marker: Final = uuid.uuid4().hex[:8]
before: Final = len(tls_peers.nonpfs.received())
response: Final = cipher_proxy.request(
"POST",
"/v1/chat/completions",
{
"model": "tls-nonpfs-peer",
"messages": [{"role": "user", "content": "tls"}],
"user": f"user15-{marker}",
},
)
assert response.status_code >= 400, response.text
assert "ssl" in response.text.lower() or "connection error" in response.text.lower(), response.text
assert len(tls_peers.nonpfs.received()) == before
readiness: Final = cipher_proxy.client.get("/health/readiness")
assert readiness.status_code == 200, readiness.text
recovered: Final = _sync_completion(cipher_proxy, "tls-peer", f"user15-ok-{marker}")
assert recovered == f"chatcmpl-user15-ok-{marker}"

View file

@ -0,0 +1,314 @@
from __future__ import annotations
import asyncio
import os
import socket
import subprocess
import sys
import textwrap
import threading
import time
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Final
import httpx
import pytest
from pydantic import JsonValue, TypeAdapter
if TYPE_CHECKING:
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests.integration._support.client import eventually
from tests.integration._support.tls import TlsPeer, tls_peer, write_self_signed_cert
CHACHA20: Final = "ECDHE-RSA-CHACHA20-POLY1305"
AES256_GCM: Final = "AES256-GCM-SHA384"
AES128_GCM: Final = "AES128-GCM-SHA256"
CBC_SHA256: Final = "ECDHE-RSA-AES128-SHA256"
CBC_SHA384: Final = "ECDHE-RSA-AES256-SHA384"
PFS_GCM: Final = "ECDHE-RSA-AES256-GCM-SHA384"
BURST_SIZE: Final = 30
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
OVERRIDE_EXCHANGE: Final = textwrap.dedent(
"""
import sys
from litellm.llms.custom_httpx.http_handler import HTTPHandler
handler = HTTPHandler(ssl_verify=sys.argv[2])
try:
response = handler.post(f"{sys.argv[1]}/v1/chat/completions", json={"user": "u-override"})
print(response.json()["id"])
finally:
handler.close()
"""
)
@pytest.fixture(scope="module")
def tls_cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]:
return write_self_signed_cert(tmp_path_factory.mktemp("tlscerts"))
@dataclass(frozen=True, slots=True)
class BurstOutcome:
user: str
when: float
response_id: str | None
error: Exception | None
def _sync_post(cert_file: Path, url: str, user: str) -> str:
from litellm.llms.custom_httpx.http_handler import HTTPHandler
handler: Final = HTTPHandler(ssl_verify=str(cert_file))
try:
response: Final = handler.client.post(f"{url}/v1/chat/completions", json={"user": user})
return str(JSON_OBJECT.validate_json(response.content)["id"])
finally:
handler.close()
def _async_post(cert_file: Path, url: str, user: str) -> str:
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
async def _run() -> str:
handler: Final = AsyncHTTPHandler(ssl_verify=str(cert_file))
try:
response: Final = await handler.client.post(f"{url}/v1/chat/completions", json={"user": user})
return str(JSON_OBJECT.validate_json(response.content)["id"])
finally:
await handler.close()
return asyncio.run(_run())
def _assert_served(peer: TlsPeer, user: str, response_id: str) -> None:
records: Final = peer.received()
assert response_id == f"chatcmpl-{user}"
assert [(record.user, record.cipher) for record in records] == [(user, peer.cipher)], records
def test_sync_handler_refuses_chacha20_only_peer(tls_cert: tuple[Path, Path]) -> None:
cert_file, key_file = tls_cert
with tls_peer(CHACHA20, cert_file, key_file) as peer:
with pytest.raises(httpx.ConnectError, match="SSL"):
_sync_post(cert_file, peer.url, "u-cell1")
assert peer.received() == (), peer.received()
def test_async_httpx_handler_refuses_chacha20_only_peer(
monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path]
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
cert_file, key_file = tls_cert
with tls_peer(CHACHA20, cert_file, key_file) as peer:
with pytest.raises(httpx.ConnectError, match="SSL"):
_async_post(cert_file, peer.url, "u-cell2")
assert peer.received() == (), peer.received()
def test_async_aiohttp_handler_refuses_chacha20_only_peer(tls_cert: tuple[Path, Path]) -> None:
cert_file, key_file = tls_cert
with tls_peer(CHACHA20, cert_file, key_file) as peer:
with pytest.raises(httpx.ProtocolError, match="SSL"):
_async_post(cert_file, peer.url, "u-cell3")
assert peer.received() == (), peer.received()
def test_sync_handler_refuses_non_pfs_aes256_gcm_peer(tls_cert: tuple[Path, Path]) -> None:
cert_file, key_file = tls_cert
with tls_peer(AES256_GCM, cert_file, key_file) as peer:
with pytest.raises(httpx.ConnectError, match="SSL"):
_sync_post(cert_file, peer.url, "u-cell4")
assert peer.received() == (), peer.received()
def test_async_httpx_handler_refuses_non_pfs_aes128_gcm_peer(
monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path]
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
cert_file, key_file = tls_cert
with tls_peer(AES128_GCM, cert_file, key_file) as peer:
with pytest.raises(httpx.ConnectError, match="SSL"):
_async_post(cert_file, peer.url, "u-cell5")
assert peer.received() == (), peer.received()
def test_sync_handler_accepts_cbc_peer_by_default(tls_cert: tuple[Path, Path]) -> None:
cert_file, key_file = tls_cert
with tls_peer(CBC_SHA256, cert_file, key_file) as peer:
_assert_served(peer, "u-cell6", _sync_post(cert_file, peer.url, "u-cell6"))
def test_sync_handler_refuses_cbc_peer_in_fips_mode(
monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path]
) -> None:
monkeypatch.setenv("LITELLM_FIPS_MODE", "true")
cert_file, key_file = tls_cert
with tls_peer(CBC_SHA256, cert_file, key_file) as peer:
with pytest.raises(httpx.ConnectError, match="SSL"):
_sync_post(cert_file, peer.url, "u-cell7")
assert peer.received() == (), peer.received()
def test_async_httpx_handler_refuses_cbc_sha384_peer_in_fips_mode(
monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path]
) -> None:
monkeypatch.setenv("LITELLM_FIPS_MODE", "true")
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
cert_file, key_file = tls_cert
with tls_peer(CBC_SHA384, cert_file, key_file) as peer:
with pytest.raises(httpx.ConnectError, match="SSL"):
_async_post(cert_file, peer.url, "u-cell8")
assert peer.received() == (), peer.received()
def test_handlers_accept_pfs_gcm_peer_in_default_and_fips_modes(
monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path]
) -> None:
cert_file, key_file = tls_cert
with tls_peer(PFS_GCM, cert_file, key_file) as peer:
served: Final = _sync_post(cert_file, peer.url, "u-cell9-sync")
assert served == "chatcmpl-u-cell9-sync"
served_async: Final = _async_post(cert_file, peer.url, "u-cell9-async")
assert served_async == "chatcmpl-u-cell9-async"
monkeypatch.setenv("LITELLM_FIPS_MODE", "true")
served_fips: Final = _sync_post(cert_file, peer.url, "u-cell9-fips")
assert served_fips == "chatcmpl-u-cell9-fips"
served_async_fips: Final = _async_post(cert_file, peer.url, "u-cell9-async-fips")
assert served_async_fips == "chatcmpl-u-cell9-async-fips"
assert {record.cipher for record in peer.received()} == {PFS_GCM}
def _override_exchange(url: str, cert_file: Path, overrides: Mapping[str, str]) -> str:
completed: Final = subprocess.run(
[sys.executable, "-P", "-c", OVERRIDE_EXCHANGE, url, str(cert_file)],
env={**os.environ, **overrides},
capture_output=True,
text=True,
timeout=60,
check=False,
)
assert completed.returncode == 0, completed.stderr
return completed.stdout.strip()
def test_ssl_ciphers_env_override_keeps_chacha20(tls_cert: tuple[Path, Path]) -> None:
cert_file, key_file = tls_cert
with tls_peer(CHACHA20, cert_file, key_file) as peer:
assert _override_exchange(peer.url, cert_file, {"LITELLM_SSL_CIPHERS": CHACHA20}) == "chatcmpl-u-override"
assert [record.cipher for record in peer.received()] == [CHACHA20]
def test_ssl_ciphers_override_wins_over_fips_mode(tls_cert: tuple[Path, Path]) -> None:
cert_file, key_file = tls_cert
with tls_peer(CHACHA20, cert_file, key_file) as peer:
result: Final = _override_exchange(
peer.url, cert_file, {"LITELLM_SSL_CIPHERS": CHACHA20, "LITELLM_FIPS_MODE": "true"}
)
assert result == "chatcmpl-u-override"
assert [record.cipher for record in peer.received()] == [CHACHA20]
def test_cipher_list_follows_fips_mode_toggle_in_one_process(
monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path]
) -> None:
cert_file, key_file = tls_cert
with tls_peer(CBC_SHA256, cert_file, key_file) as peer:
assert _sync_post(cert_file, peer.url, "u-cell12-before") == "chatcmpl-u-cell12-before"
monkeypatch.setenv("LITELLM_FIPS_MODE", "true")
with pytest.raises(httpx.ConnectError, match="SSL"):
_sync_post(cert_file, peer.url, "u-cell12-fips")
monkeypatch.delenv("LITELLM_FIPS_MODE")
assert _sync_post(cert_file, peer.url, "u-cell12-after") == "chatcmpl-u-cell12-after"
assert [record.user for record in peer.received()] == ["u-cell12-before", "u-cell12-after"]
def _assert_burst_outcomes(outcomes: Mapping[int, BurstOutcome]) -> None:
for index, outcome in outcomes.items():
if outcome.error is None:
assert outcome.response_id == f"chatcmpl-user-{index}", outcome
continue
message: str = str(outcome.error).lower()
assert "handshake" not in message and "cipher" not in message, (
f"user-{index} failed with a cipher class error: {outcome.error!r}"
)
assert isinstance(outcome.error, httpx.HTTPError), f"user-{index}: {outcome.error!r}"
def _peer_ready(port: int) -> bool:
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
return True
except OSError:
return False
def test_burst_survives_peer_restart_without_cipher_errors(tls_cert: tuple[Path, Path]) -> None:
cert_file, key_file = tls_cert
with tls_peer(PFS_GCM, cert_file, key_file) as peer:
second_wave: Final = threading.Event()
outcomes: Final[dict[int, BurstOutcome]] = {}
def _one(index: int) -> None:
if index >= 10:
second_wave.wait(timeout=30)
try:
response_id: Final = _sync_post(cert_file, peer.url, f"user-{index}")
outcomes[index] = BurstOutcome(f"user-{index}", time.monotonic(), response_id, None)
except Exception as error:
outcomes[index] = BurstOutcome(f"user-{index}", time.monotonic(), None, error)
async def _async_burst() -> None:
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
handler: Final = AsyncHTTPHandler(ssl_verify=str(cert_file))
try:
await asyncio.to_thread(second_wave.wait, 30)
await asyncio.gather(*(_async_one(handler, index) for index in range(15, BURST_SIZE)))
finally:
await handler.close()
async def _async_one(handler: AsyncHTTPHandler, index: int) -> None:
try:
response: Final = await handler.client.post(
f"{peer.url}/v1/chat/completions", json={"user": f"user-{index}"}
)
outcomes[index] = BurstOutcome(
f"user-{index}", time.monotonic(), str(JSON_OBJECT.validate_json(response.content)["id"]), None
)
except Exception as error:
outcomes[index] = BurstOutcome(f"user-{index}", time.monotonic(), None, error)
def _async_thread() -> None:
asyncio.run(_async_burst())
async_thread: Final = threading.Thread(target=_async_thread, daemon=True)
with ThreadPoolExecutor(max_workers=15) as pool:
async_thread.start()
futures: Final = tuple(pool.submit(_one, index) for index in range(15))
eventually(lambda: len(peer.received()) >= 3, lambda ready: ready, seconds=15)
stopped_at: Final = time.monotonic()
peer.stop()
peer.start()
eventually(lambda: _peer_ready(peer.port), lambda ready: ready, seconds=15)
restarted_at: Final = time.monotonic()
second_wave.set()
for future in futures:
future.result(timeout=60)
async_thread.join(timeout=60)
assert len(outcomes) == BURST_SIZE
_assert_burst_outcomes(outcomes)
succeeded: Final = {o.user for o in outcomes.values() if o.error is None}
recorded: Final = peer.received()
assert sorted(record.user or "" for record in recorded) == sorted(succeeded)
assert all(record.cipher == PFS_GCM for record in recorded)
for user in succeeded:
assert sum(1 for record in recorded if record.user == user) == 1, user
assert any(o.error is None and o.when < stopped_at for o in outcomes.values()), "no success before restart"
assert any(o.error is None and o.when > restarted_at for o in outcomes.values()), "no success after restart"