fastapi/tests/test_pydantic_json_response_class.py

import math
from typing import List, Optional

import pytest
from dirty_equals import IsFloatNan
from fastapi import FastAPI
from fastapi._compat import PYDANTIC_V2
from fastapi.responses import PydanticJSONResponse
from fastapi.testclient import TestClient
from pydantic import BaseModel, Field

from .utils import needs_py_lt_314, needs_pydanticv2

app = FastAPI(default_response_class=PydanticJSONResponse)


class CustomResponse(PydanticJSONResponse):
    media_type = "application/x-custom"


class Item(BaseModel):
    name: str
    price: float
    category: str = Field("food", alias="CAT")
    tax: float = 8.875
    description: Optional[str] = None


@app.get("/response-model", response_model=Item)
@app.get(
    "/response-model-include",
    response_model=Item,
    response_model_include={"name", "category"},
)
@app.get(
    "/response-model-exclude",
    response_model=Item,
    response_model_exclude={"tax", "description"},
)
@app.get(
    "/response-model-by-alias-false",
    response_model=Item,
    response_model_by_alias=False,
)
@app.get(
    "/response-model-exclude-unset",
    response_model=Item,
    response_model_exclude_unset=True,
)
@app.get(
    "/response-model-exclude-defaults",
    response_model=Item,
    response_model_exclude_defaults=True,
)
@app.get(
    "/response-model-exclude-none",
    response_model=Item,
    response_model_exclude_none=True,
)
def get_response_model_params():
    return {"name": "cheese", "price": 1.23, "tax": 8.875, "description": None}


class FloatsNone(BaseModel):
    # pydantic converts inf/nan to None by default
    numbers: List[float]


class FloatsNum(FloatsNone):
    model_config = {"ser_json_inf_nan": "constants"}


class FloatsStr(FloatsNone):
    model_config = {"ser_json_inf_nan": "strings"}


@app.get("/floats-none", response_model=FloatsNone)
@app.get("/floats-num", response_model=FloatsNum)
@app.get("/floats-str", response_model=FloatsStr)
@app.get("/custom-class", response_class=CustomResponse, response_model=FloatsStr)
def get_floats():
    return {"numbers": [3.14, math.inf, math.nan, 2.72]}


client = TestClient(app)


@needs_pydanticv2
@pytest.mark.parametrize(
    "path,expected_response",
    [
        (
            "/response-model",
            {
                "name": "cheese",
                "price": 1.23,
                "CAT": "food",
                "tax": 8.875,
                "description": None,
            },
        ),
        ("/response-model-include", {"name": "cheese", "CAT": "food"}),
        ("/response-model-exclude", {"name": "cheese", "price": 1.23, "CAT": "food"}),
        (
            "/response-model-by-alias-false",
            {
                "name": "cheese",
                "price": 1.23,
                "category": "food",
                "tax": 8.875,
                "description": None,
            },
        ),
        (
            "/response-model-exclude-unset",
            {
                "name": "cheese",
                "price": 1.23,
                "tax": 8.875,
                "description": None,
            },
        ),
        ("/response-model-exclude-defaults", {"name": "cheese", "price": 1.23}),
        (
            "/response-model-exclude-none",
            {
                "name": "cheese",
                "price": 1.23,
                "CAT": "food",
                "tax": 8.875,
            },
        ),
    ],
)
def test_response_model_params(path: str, expected_response: dict):
    response = client.get(path)
    assert response.status_code == 200
    assert response.json() == expected_response


@needs_pydanticv2
@pytest.mark.parametrize(
    "path,expected_numbers",
    [
        ("/floats-none", [3.14, None, None, 2.72]),
        ("/floats-num", [3.14, math.inf, IsFloatNan, 2.72]),
        ("/floats-str", [3.14, "Infinity", "NaN", 2.72]),
    ],
)
def test_floats(path: str, expected_numbers: list):
    response = client.get(path)
    assert response.status_code == 200
    assert response.json() == {"numbers": expected_numbers}


@needs_pydanticv2
def test_custom_response_class():
    response = client.get("/custom-class")
    assert response.status_code == 200
    assert response.headers["content-type"] == "application/x-custom"
    assert response.json() == {"numbers": [3.14, "Infinity", "NaN", 2.72]}


@needs_py_lt_314
def test_requires_pydantic_v2_model():
    if PYDANTIC_V2:
        from pydantic.v1 import BaseModel as BaseModelV1
    else:
        from pydantic import BaseModel as BaseModelV1

    app = FastAPI(default_response_class=PydanticJSONResponse)

    class ModelV1(BaseModelV1):
        data: str

    @app.get("/model-v1")
    def get_model_v1() -> ModelV1:
        return ModelV1(data="hello")

    client = TestClient(app)
    with pytest.raises(AssertionError, match="requires a pydantic v2 model"):
        client.get("/model-v1")