tests for pydantic generic model support

mike-oakley · Dec 3, 2024 · 9575fa7 · 9575fa7
1 parent fb3cc7f
commit 9575fa7
Show file tree

Hide file tree

Showing 3 changed files with 155 additions and 3 deletions.
diff --git a/openapi_pydantic/v3/v3_0/util.py b/openapi_pydantic/v3/v3_0/util.py
@@ -1,4 +1,5 @@
 import logging
+import re
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -242,6 +243,9 @@ def _traverse(obj: Any) -> None:
 
 
 def _construct_ref_obj(pydantic_schema: PydanticSchema[PydanticType]) -> Reference:
-    ref_obj = Reference(**{"$ref": ref_prefix + pydantic_schema.schema_class.__name__})
+    ref_name = re.sub(
+        r"[^a-zA-Z0-9.\-_]", "_", pydantic_schema.schema_class.__name__
+    ).replace(".", "__")
+    ref_obj = Reference(**{"$ref": ref_prefix + ref_name})
     logger.debug(f"ref_obj={ref_obj}")
     return ref_obj
diff --git a/tests/util/test_util.py b/tests/util/test_util.py
@@ -1,6 +1,7 @@
 import logging
-from typing import Callable
+from typing import Callable, Generic, TypeVar
 
+import pytest
 from pydantic import BaseModel, Field
 
 from openapi_pydantic import (
@@ -68,6 +69,43 @@ def test_construct_open_api_with_schema_class_3() -> None:
     assert "resp_bar" in schema_without_alias.properties
 
 
+@pytest.mark.skipif(PYDANTIC_V2, reason="generic type for Pydantic V1")
+def test_construct_open_api_with_schema_class_4_generic_response_v1() -> None:
+    DataT = TypeVar("DataT")
+    from pydantic.v1.generics import GenericModel
+
+    class GenericResponse(GenericModel, Generic[DataT]):
+        msg: str = Field(description="message of the generic response")
+        data: DataT = Field(description="data value of the generic response")
+
+    open_api_4 = construct_base_open_api_4_generic_response(
+        GenericResponse[PongResponse]
+    )
+
+    result = construct_open_api_with_schema_class(open_api_4)
+    assert result.components is not None
+    assert result.components.schemas is not None
+    assert "GenericResponse_PongResponse_" in result.components.schemas
+
+
+@pytest.mark.skipif(not PYDANTIC_V2, reason="generic type for Pydantic V2")
+def test_construct_open_api_with_schema_class_4_generic_response_v2() -> None:
+    DataT = TypeVar("DataT")
+
+    class GenericResponse(BaseModel, Generic[DataT]):
+        msg: str = Field(description="message of the generic response")
+        data: DataT = Field(description="data value of the generic response")
+
+    open_api_4 = construct_base_open_api_4_generic_response(
+        GenericResponse[PongResponse]
+    )
+
+    result = construct_open_api_with_schema_class(open_api_4)
+    assert result.components is not None
+    assert result.components.schemas is not None
+    assert "GenericResponse_PongResponse_" in result.components.schemas
+
+
 def construct_base_open_api_1() -> OpenAPI:
     model_validate: Callable[[dict], OpenAPI] = getattr(
         OpenAPI, "model_validate" if PYDANTIC_V2 else "parse_obj"
@@ -176,6 +214,42 @@ def construct_base_open_api_3() -> OpenAPI:
     )
 
 
+def construct_base_open_api_4_generic_response(response_schema: type) -> OpenAPI:
+    return OpenAPI(
+        info=Info(
+            title="My own API",
+            version="v0.0.1",
+        ),
+        paths={
+            "/ping": PathItem(
+                post=Operation(
+                    requestBody=RequestBody(
+                        content={
+                            "application/json": MediaType(
+                                media_type_schema=PydanticSchema(
+                                    schema_class=PingRequest
+                                )
+                            )
+                        }
+                    ),
+                    responses={
+                        "200": Response(
+                            description="pong",
+                            content={
+                                "application/json": MediaType(
+                                    media_type_schema=PydanticSchema(
+                                        schema_class=response_schema
+                                    )
+                                )
+                            },
+                        )
+                    },
+                )
+            )
+        },
+    )
+
+
 class PingRequest(BaseModel):
     """Ping Request"""
 

diff --git a/tests/v3_0/test_util.py b/tests/v3_0/test_util.py
@@ -1,6 +1,7 @@
 import logging
-from typing import Callable, Literal
+from typing import Callable, Generic, Literal, TypeVar
 
+import pytest
 from pydantic import BaseModel, Field
 
 from openapi_pydantic.compat import PYDANTIC_V2
@@ -74,6 +75,43 @@ def test_construct_open_api_with_schema_class_3() -> None:
     assert "resp_bar" in schema_without_alias.properties
 
 
+@pytest.mark.skipif(PYDANTIC_V2, reason="generic type for Pydantic V1")
+def test_construct_open_api_with_schema_class_4_generic_response_v1() -> None:
+    DataT = TypeVar("DataT")
+    from pydantic.v1.generics import GenericModel
+
+    class GenericResponse(GenericModel, Generic[DataT]):
+        msg: str = Field(description="message of the generic response")
+        data: DataT = Field(description="data value of the generic response")
+
+    open_api_4 = construct_base_open_api_4_generic_response(
+        GenericResponse[PongResponse]
+    )
+
+    result = construct_open_api_with_schema_class(open_api_4)
+    assert result.components is not None
+    assert result.components.schemas is not None
+    assert "GenericResponse_PongResponse_" in result.components.schemas
+
+
+@pytest.mark.skipif(not PYDANTIC_V2, reason="generic type for Pydantic V2")
+def test_construct_open_api_with_schema_class_4_generic_response() -> None:
+    DataT = TypeVar("DataT")
+
+    class GenericResponse(BaseModel, Generic[DataT]):
+        msg: str = Field(description="message of the generic response")
+        data: DataT = Field(description="data value of the generic response")
+
+    open_api_4 = construct_base_open_api_4_generic_response(
+        GenericResponse[PongResponse]
+    )
+
+    result = construct_open_api_with_schema_class(open_api_4)
+    assert result.components is not None
+    assert result.components.schemas is not None
+    assert "GenericResponse_PongResponse_" in result.components.schemas
+
+
 def construct_base_open_api_1() -> OpenAPI:
     model_validate: Callable[[dict], OpenAPI] = getattr(
         OpenAPI, "model_validate" if PYDANTIC_V2 else "parse_obj"
@@ -215,6 +253,42 @@ def construct_base_open_api_3_plus() -> OpenAPI:
     )
 
 
+def construct_base_open_api_4_generic_response(response_schema: type) -> OpenAPI:
+    return OpenAPI(
+        info=Info(
+            title="My own API",
+            version="v0.0.1",
+        ),
+        paths={
+            "/ping": PathItem(
+                post=Operation(
+                    requestBody=RequestBody(
+                        content={
+                            "application/json": MediaType(
+                                media_type_schema=PydanticSchema(
+                                    schema_class=PingRequest
+                                )
+                            )
+                        }
+                    ),
+                    responses={
+                        "200": Response(
+                            description="pong",
+                            content={
+                                "application/json": MediaType(
+                                    media_type_schema=PydanticSchema(
+                                        schema_class=response_schema
+                                    )
+                                )
+                            },
+                        )
+                    },
+                )
+            )
+        },
+    )
+
+
 class PingRequest(BaseModel):
     """Ping Request"""