119 lines
4.2 KiB
Python
119 lines
4.2 KiB
Python
import json
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from ai_platform.infrastructure.aihub import AiHubClient, AiHubError
|
|
|
|
VKEY = "request-only-vkey-test"
|
|
|
|
|
|
def client_for(handler, *, timeout: float = 0.2) -> tuple[AiHubClient, httpx.AsyncClient]:
|
|
transport = httpx.MockTransport(handler)
|
|
http_client = httpx.AsyncClient(transport=transport)
|
|
return AiHubClient("https://aihub.test", http_client, timeout), http_client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sends_scoring_request_and_parses_scored_json():
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "POST"
|
|
assert request.url.path == "/v1/chat/completions"
|
|
assert request.headers["authorization"] == f"Bearer {VKEY}"
|
|
assert "x-tenant-vkey" not in request.headers
|
|
body = json.loads(request.content)
|
|
assert body["stream"] is False
|
|
assert body["response_format"] == {"type": "json_object"}
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": '{"score": 82, "confidence": 0.93, "reasoning": "Strong fit"}'
|
|
}
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
adapter, http_client = client_for(handler)
|
|
try:
|
|
result = await adapter.score(
|
|
model="catalog-model", system_prompt="System", user_prompt="Snapshot", vkey=VKEY
|
|
)
|
|
assert result.score == 82
|
|
assert result.confidence == 0.93
|
|
assert result.reasoning == "Strong fit"
|
|
finally:
|
|
await http_client.aclose()
|
|
|
|
|
|
@pytest.mark.parametrize("status", [401, 402, 429, 500, 503])
|
|
@pytest.mark.asyncio
|
|
async def test_http_errors_are_sanitized(status: int):
|
|
adapter, http_client = client_for(lambda _request: httpx.Response(status, text=VKEY))
|
|
try:
|
|
with pytest.raises(AiHubError, match=f"HTTP {status}") as error:
|
|
await adapter.score(model="m", system_prompt="s", user_prompt="u", vkey=VKEY)
|
|
assert VKEY not in str(error.value)
|
|
assert error.value.error_code == "aihub_http_error"
|
|
finally:
|
|
await http_client.aclose()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"body",
|
|
[
|
|
b"not-json",
|
|
b'{"choices": []}',
|
|
b'{"choices": [{"message": {"content": "not-json"}}]}',
|
|
b'{"choices": [{"message": {"content": "{\\"score\\": 5}"}}]}',
|
|
b'{"choices": [{"message": {"content": "{\\"score\\": 5, \\"score\\": 6, \\"confidence\\": 0.5, \\"reasoning\\": \\"duplicate\\"}"}}]}',
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_malformed_scoring_responses_are_rejected(body: bytes):
|
|
adapter, http_client = client_for(lambda _request: httpx.Response(200, content=body))
|
|
try:
|
|
with pytest.raises(AiHubError, match="invalid JSON response"):
|
|
await adapter.score(model="m", system_prompt="s", user_prompt="u", vkey=VKEY)
|
|
finally:
|
|
await http_client.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_timeout_is_sanitized():
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
raise httpx.ReadTimeout("upstream error includes " + VKEY)
|
|
|
|
adapter, http_client = client_for(handler)
|
|
try:
|
|
with pytest.raises(AiHubError, match="timed out") as error:
|
|
await adapter.score(model="m", system_prompt="s", user_prompt="u", vkey=VKEY)
|
|
assert VKEY not in str(error.value)
|
|
assert error.value.error_code == "aihub_timeout"
|
|
finally:
|
|
await http_client.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_response_echoing_vkey_is_rejected_without_exposing_it():
|
|
response = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": json.dumps({"score": 70, "confidence": 0.8, "reasoning": VKEY})
|
|
}
|
|
}
|
|
]
|
|
}
|
|
adapter, http_client = client_for(lambda _request: httpx.Response(200, json=response))
|
|
try:
|
|
with pytest.raises(AiHubError) as error:
|
|
await adapter.score(model="m", system_prompt="s", user_prompt="u", vkey=VKEY)
|
|
assert VKEY not in str(error.value)
|
|
assert error.value.error_code == "aihub_invalid_response"
|
|
finally:
|
|
await http_client.aclose()
|