80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
from typing import List
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from great_ai import GreatAI, parameter
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class Score(BaseModel):
|
|
label: str
|
|
confidence: float
|
|
|
|
|
|
@pytest.mark.parametrize("is_async", [False, True])
|
|
def test_custom_validation_returns_client_error_but_prediction_failures_do_not(
|
|
isolated_context: None, is_async: bool
|
|
) -> None:
|
|
@parameter("value", validate=lambda value: value > 0)
|
|
def predict(value: int) -> int:
|
|
if value == 7:
|
|
raise ValueError("prediction failed")
|
|
return value
|
|
|
|
@parameter("value", validate=lambda value: value > 0)
|
|
async def predict_async(value: int) -> int:
|
|
return predict(value)
|
|
|
|
service = GreatAI.create(predict_async if is_async else predict)
|
|
with TestClient(service.app) as client:
|
|
rejected = client.post("/predict", json={"value": -1})
|
|
assert rejected.status_code == 422
|
|
assert "did not pass validation" in rejected.json()["detail"]
|
|
assert client.post("/predict", json={"value": "wrong"}).status_code == 422
|
|
assert client.post("/predict", json={"value": 7}).status_code == 500
|
|
valid = client.post("/predict", json={"value": 1})
|
|
assert valid.status_code == 200
|
|
assert valid.json()["output"] == 1
|
|
|
|
|
|
@pytest.mark.parametrize("is_async", [False, True])
|
|
def test_openapi_preserves_typed_nested_prediction_outputs(
|
|
isolated_context: None, is_async: bool
|
|
) -> None:
|
|
def predict(text: str) -> List[Score]:
|
|
return [Score(label=text, confidence=0.75)]
|
|
|
|
async def predict_async(text: str) -> List[Score]:
|
|
return predict(text)
|
|
|
|
service = GreatAI.create(predict_async if is_async else predict)
|
|
with TestClient(service.app) as client:
|
|
schema = client.get("/openapi.json").json()
|
|
response_ref = schema["paths"]["/predict"]["post"]["responses"]["200"][
|
|
"content"
|
|
]["application/json"]["schema"]["$ref"].split("/")[-1]
|
|
output = schema["components"]["schemas"][response_ref]["properties"]["output"][
|
|
"anyOf"
|
|
][0]
|
|
assert output == {
|
|
"items": {"$ref": "#/components/schemas/Score"},
|
|
"type": "array",
|
|
}
|
|
assert (
|
|
schema["components"]["schemas"]["Score"]["properties"]["confidence"]["type"]
|
|
== "number"
|
|
)
|
|
response = client.post("/predict", json={"text": "cat"})
|
|
assert response.status_code == 200
|
|
assert response.json()["output"] == [{"label": "cat", "confidence": 0.75}]
|
|
|
|
|
|
def test_http_rejects_outputs_that_violate_the_return_annotation(
|
|
isolated_context: None,
|
|
) -> None:
|
|
@GreatAI.create
|
|
def predict(value: int) -> int:
|
|
return "invalid" # type: ignore[return-value]
|
|
|
|
with TestClient(predict.app, raise_server_exceptions=False) as client:
|
|
assert client.post("/predict", json={"value": 1}).status_code == 500
|