Files
jundot--omlx/tests/test_external_api.py
wehub-resource-sync e9a2f726c9
CI / test (3.11) (push) Has been cancelled
CI / test (3.12) (push) Has been cancelled
CI / test (3.13) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:29:51 +08:00

460 lines
15 KiB
Python

# SPDX-License-Identifier: Apache-2.0
"""Tests for the external OpenAI-compatible endpoint client."""
import json
import httpx
import pytest
from omlx.admin.external_api import (
ExternalAPIClient,
ExternalChatAdapter,
ExternalEndpointConfig,
ExternalEndpointError,
)
API_KEY = "sk-secret-test-key"
def _config(**overrides):
params = {
"base_url": "https://api.example.com/v1",
"api_key": API_KEY,
"model": "test-model",
}
params.update(overrides)
return ExternalEndpointConfig(**params)
def _client(handler, **config_overrides):
return ExternalAPIClient(
_config(**config_overrides),
transport=httpx.MockTransport(handler),
)
def _sse_body(chunks, done=True):
"""Build an SSE response body from chunk dicts / raw lines."""
lines = []
for chunk in chunks:
if isinstance(chunk, str):
lines.append(chunk)
else:
lines.append("data: " + json.dumps(chunk))
if done:
lines.append("data: [DONE]")
return ("\n\n".join(lines) + "\n\n").encode()
def _content_chunk(text):
return {"choices": [{"delta": {"content": text}}]}
def _usage_chunk(prompt=10, completion=5, cached=None):
usage = {"prompt_tokens": prompt, "completion_tokens": completion}
if cached is not None:
usage["prompt_tokens_details"] = {"cached_tokens": cached}
return {"choices": [], "usage": usage}
def _completion_response(text="A", prompt=10, completion=5, extra_message=None):
message = {"role": "assistant", "content": text}
if extra_message:
message.update(extra_message)
return httpx.Response(
200,
json={
"choices": [{"message": message, "finish_reason": "stop"}],
"usage": {"prompt_tokens": prompt, "completion_tokens": completion},
},
)
# =============================================================================
# Config validation
# =============================================================================
class TestExternalEndpointConfig:
def test_trailing_slash_stripped(self):
cfg = _config(base_url="https://api.example.com/v1/")
assert cfg.base_url == "https://api.example.com/v1"
def test_whitespace_stripped(self):
cfg = _config(base_url=" http://localhost:8001/v1 ", model=" m ")
assert cfg.base_url == "http://localhost:8001/v1"
assert cfg.model == "m"
def test_invalid_scheme_rejected(self):
with pytest.raises(ValueError, match="http:// or https://"):
_config(base_url="ftp://api.example.com/v1")
def test_empty_model_rejected(self):
with pytest.raises(ValueError, match="model must not be empty"):
_config(model=" ")
def test_api_key_hidden_from_repr(self):
cfg = _config()
assert API_KEY not in repr(cfg)
assert API_KEY not in str(cfg)
assert cfg.api_key.get_secret_value() == API_KEY
def test_empty_api_key_allowed(self):
cfg = _config(api_key="")
assert cfg.api_key.get_secret_value() == ""
# =============================================================================
# Streaming chat completion
# =============================================================================
class TestStreamChatCompletion:
async def test_counts_come_from_usage_not_chunks(self):
"""Providers batch multiple tokens per chunk — only usage counts."""
captured = {}
def handler(request):
captured["body"] = json.loads(request.content)
captured["auth"] = request.headers.get("authorization")
return httpx.Response(
200,
content=_sse_body([
{"choices": [{"delta": {"role": "assistant"}}]},
_content_chunk("Hello "),
_content_chunk("world"),
_usage_chunk(prompt=1000, completion=128),
]),
)
client = _client(handler)
try:
stats = await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=128,
temperature=0.0,
)
finally:
await client.aclose()
# 2 content chunks but usage says 128 completion tokens
assert stats.completion_tokens == 128
assert stats.prompt_tokens == 1000
assert stats.cached_tokens == 0
assert stats.text == "Hello world"
assert stats.start_time <= stats.first_content_time
assert stats.first_content_time <= stats.last_content_time <= stats.end_time
assert captured["auth"] == f"Bearer {API_KEY}"
assert captured["body"]["stream"] is True
assert captured["body"]["stream_options"] == {"include_usage": True}
async def test_missing_usage_is_hard_error(self):
def handler(request):
return httpx.Response(
200,
content=_sse_body([_content_chunk("Hello")]),
)
client = _client(handler)
try:
with pytest.raises(ExternalEndpointError, match="stream usage"):
await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
async def test_cached_tokens_parsed(self):
def handler(request):
return httpx.Response(
200,
content=_sse_body([
_content_chunk("x"),
_usage_chunk(prompt=100, completion=10, cached=64),
]),
)
client = _client(handler)
try:
stats = await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
assert stats.cached_tokens == 64
async def test_keepalive_and_malformed_lines_skipped(self):
def handler(request):
return httpx.Response(
200,
content=_sse_body([
": keepalive",
"data: not-json",
_content_chunk("ok"),
_usage_chunk(),
]),
)
client = _client(handler)
try:
stats = await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
assert stats.text == "ok"
async def test_single_chunk_dump_has_zero_gen_duration(self):
def handler(request):
return httpx.Response(
200,
content=_sse_body([
_content_chunk("everything at once"),
_usage_chunk(prompt=10, completion=20),
]),
)
client = _client(handler)
try:
stats = await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=32,
temperature=0.0,
)
finally:
await client.aclose()
assert stats.first_content_time == stats.last_content_time
async def test_reasoning_content_sets_timing_but_not_text(self):
def handler(request):
return httpx.Response(
200,
content=_sse_body([
{"choices": [{"delta": {"reasoning_content": "thinking..."}}]},
_content_chunk("answer"),
_usage_chunk(),
]),
)
client = _client(handler)
try:
stats = await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
assert stats.text == "answer"
assert stats.first_content_time < stats.last_content_time
async def test_auth_error_message_without_key(self):
def handler(request):
return httpx.Response(401, json={"error": {"message": "bad key"}})
client = _client(handler)
try:
with pytest.raises(ExternalEndpointError) as exc_info:
await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
assert "rejected the API key" in str(exc_info.value)
assert API_KEY not in str(exc_info.value)
async def test_http_error_includes_json_detail(self):
def handler(request):
return httpx.Response(
500, json={"error": {"message": "model exploded"}}
)
client = _client(handler)
try:
with pytest.raises(ExternalEndpointError, match="model exploded"):
await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
async def test_connect_error_mapped(self):
def handler(request):
raise httpx.ConnectError("connection refused", request=request)
client = _client(handler)
try:
with pytest.raises(ExternalEndpointError, match="Cannot connect"):
await client.stream_chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
# =============================================================================
# Non-streaming chat completion
# =============================================================================
class TestChatCompletion:
async def test_returns_text_and_usage(self):
def handler(request):
return _completion_response(text="B", prompt=42, completion=7)
client = _client(handler)
try:
result = await client.chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
assert result.text == "B"
assert result.prompt_tokens == 42
assert result.completion_tokens == 7
async def test_unexpected_shape_raises(self):
def handler(request):
return httpx.Response(200, json={"unexpected": True})
client = _client(handler)
try:
with pytest.raises(ExternalEndpointError, match="unexpected response"):
await client.chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=0.0,
)
finally:
await client.aclose()
async def test_no_auth_header_without_key(self):
captured = {}
def handler(request):
captured["auth"] = request.headers.get("authorization")
return _completion_response()
client = _client(handler, api_key="")
try:
await client.chat_completion(
messages=[{"role": "user", "content": "hi"}],
max_tokens=8,
temperature=None,
)
finally:
await client.aclose()
assert captured["auth"] is None
# =============================================================================
# Accuracy adapter
# =============================================================================
class TestExternalChatAdapter:
def _adapter(self, handler, profile="deterministic"):
client = _client(handler)
return ExternalChatAdapter(client, profile), client
async def test_deterministic_sends_temperature_zero_whitelist_only(self):
captured = {}
def handler(request):
captured["body"] = json.loads(request.content)
return _completion_response()
adapter, client = self._adapter(handler)
try:
# Simulate _eval_single's setdefault-injected kwargs — the
# adapter must drop everything except the whitelist.
output = await adapter.chat(
messages=[{"role": "user", "content": "q"}],
max_tokens=128,
temperature=0.7,
presence_penalty=0.0,
repetition_penalty=1.0,
chat_template_kwargs={"enable_thinking": False},
)
finally:
await client.aclose()
assert output.text == "A"
body = captured["body"]
assert set(body.keys()) == {"model", "messages", "max_tokens", "temperature"}
assert body["temperature"] == 0.0
assert body["max_tokens"] == 128
async def test_model_settings_omits_sampling_params(self):
captured = {}
def handler(request):
captured["body"] = json.loads(request.content)
return _completion_response()
adapter, client = self._adapter(handler, profile="model_settings")
try:
await adapter.chat(
messages=[{"role": "user", "content": "q"}],
max_tokens=64,
temperature=0.0,
)
finally:
await client.aclose()
body = captured["body"]
assert set(body.keys()) == {"model", "messages", "max_tokens"}
async def test_text_uses_content_not_reasoning(self):
def handler(request):
return _completion_response(
text="final", extra_message={"reasoning_content": "hidden"}
)
adapter, client = self._adapter(handler)
try:
output = await adapter.chat(messages=[{"role": "user", "content": "q"}])
finally:
await client.aclose()
assert output.text == "final"
def test_model_type_is_none(self):
def handler(request):
return _completion_response()
adapter, _ = self._adapter(handler)
assert adapter.model_type is None
async def test_preflight_raises_on_auth_error(self):
def handler(request):
return httpx.Response(401, json={"error": {"message": "nope"}})
adapter, client = self._adapter(handler)
try:
with pytest.raises(ExternalEndpointError, match="rejected the API key"):
await adapter.preflight()
finally:
await client.aclose()
async def test_preflight_passes_on_success(self):
def handler(request):
return _completion_response(text="OK")
adapter, client = self._adapter(handler)
try:
await adapter.preflight()
finally:
await client.aclose()