Repository: airtai/faststream
Pull request #2150 (2025-03-28 17:58:41)

fix: pydantic211 compat


Patch:
diff --git a/faststream/asyncapi/schema/main.py b/faststream/asyncapi/schema/main.py
index fc11c108..d90f0425 100644
--- a/faststream/asyncapi/schema/main.py
+++ b/faststream/asyncapi/schema/main.py
@@ -9,9 +9,7 @@ from faststream.asyncapi.schema.message import Message
 from faststream.asyncapi.schema.servers import Server
 from faststream.asyncapi.schema.utils import (
     ExternalDocs,
-    ExternalDocsDict,
     Tag,
-    TagDict,
 )
 from faststream.exceptions import INSTALL_YAML
 
@@ -95,8 +93,8 @@ class Schema(BaseModel):
     servers: Optional[Dict[str, Server]] = None
     channels: Dict[str, Channel]
     components: Optional[Components] = None
-    tags: Optional[List[Union[Tag, TagDict, Dict[str, Any]]]] = None
-    externalDocs: Optional[Union[ExternalDocs, ExternalDocsDict, Dict[str, Any]]] = None
+    tags: Optional[List[Union[Tag, Dict[str, Any]]]] = None
+    externalDocs: Optional[Union[ExternalDocs, Dict[str, Any]]] = None
 
     def to_jsonable(self) -> Any:
         """Convert the schema to a JSON-serializable object."""


Test patch:
diff --git a/tests/asyncapi/base/arguments.py b/tests/asyncapi/base/arguments.py
index 3a17bbce..1af45d65 100644
--- a/tests/asyncapi/base/arguments.py
+++ b/tests/asyncapi/base/arguments.py
@@ -16,6 +16,8 @@ from tests.marks import pydantic_v2
 
 
 class FastAPICompatible:
+    is_fastapi: bool = False
+
     broker_class: Type[BrokerUsecase]
     dependency_builder = staticmethod(APIDepends)
 
@@ -508,35 +510,45 @@ class FastAPICompatible:
         class Sub(pydantic.BaseModel):
             type: Literal["sub"]
 
-        descriminator = Annotated[
-            Union[Sub2, Sub], pydantic.Field(discriminator="type")
-        ]
-
         broker = self.broker_class()
 
         @broker.subscriber("test")
-        async def handle(user: descriminator): ...
+        async def handle(
+            user: Annotated[Union[Sub2, Sub], pydantic.Field(discriminator="type")],
+        ): ...
 
         schema = get_app_schema(self.build_app(broker)).to_jsonable()
 
         key = next(iter(schema["components"]["messages"].keys()))
+
         assert key == IsStr(regex=r"test[\w:]*:Handle:Message")
-        assert schema["components"] == {
-            "messages": {
-                key: {
-                    "title": key,
-                    "correlationId": {"location": "$message.header#/correlation_id"},
-                    "payload": {
-                        "discriminator": "type",
-                        "oneOf": [
-                            {"$ref": "#/components/schemas/Sub2"},
-                            {"$ref": "#/components/schemas/Sub"},
-                        ],
-                        "title": "Handle:Message:Payload",
-                    },
-                }
-            },
-            "schemas": {
+
+        expected_schema = IsPartialDict(
+            {
+                "discriminator": "type",
+                "oneOf": [
+                    {"$ref": "#/components/schemas/Sub2"},
+                    {"$ref": "#/components/schemas/Sub"},
+                ],
+                "title": "Handle:Message:Payload",
+            }
+        )
+        if self.is_fastapi:
+            expected_schema = (
+                IsPartialDict(
+                    {
+                        "$ref": "#/components/schemas/Handle:Message:Payload",
+                    }
+                )
+                | expected_schema
+            )
+
+        assert schema["components"]["messages"][key]["payload"] == expected_schema, (
+            schema["components"]
+        )
+
+        assert schema["components"]["schemas"] == IsPartialDict(
+            {
                 "Sub": {
                     "properties": {
                         "type": IsPartialDict({"const": "sub", "title": "Type"})
@@ -553,8 +565,20 @@ class FastAPICompatible:
                     "title": "Sub2",
                     "type": "object",
                 },
-            },
-        }, schema["components"]
+            }
+        ), schema["components"]["schemas"]
+
+        if self.is_fastapi and (
+            payload := schema["components"]["schemas"].get("Handle:Message:Payload")
+        ):
+            assert payload == IsPartialDict(
+                {
+                    "anyOf": [
+                        {"$ref": "#/components/schemas/Sub2"},
+                        {"$ref": "#/components/schemas/Sub"},
+                    ]
+                }
+            )
 
     @pydantic_v2
     def test_nested_descriminator(self):
@@ -578,11 +602,11 @@ class FastAPICompatible:
         assert key == IsStr(regex=r"test[\w:]*:Handle:Message")
         assert schema["components"] == {
             "messages": {
-                key: {
-                    "title": key,
-                    "correlationId": {"location": "$message.header#/correlation_id"},
-                    "payload": {"$ref": "#/components/schemas/Model"},
-                }
+                key: IsPartialDict(
+                    {
+                        "payload": {"$ref": "#/components/schemas/Model"},
+                    }
+                )
             },
             "schemas": {
                 "Sub": {
diff --git a/tests/asyncapi/base/fastapi.py b/tests/asyncapi/base/fastapi.py
index 58ded715..6f5170e3 100644
--- a/tests/asyncapi/base/fastapi.py
+++ b/tests/asyncapi/base/fastapi.py
@@ -12,6 +12,8 @@ from faststream.broker.types import MsgType
 
 
 class FastAPITestCase:
+    is_fastapi = True
+
     broker_class: Type[StreamRouter[MsgType]]
     broker_wrapper: Callable[[BrokerUsecase[MsgType, Any]], BrokerUsecase[MsgType, Any]]
 
diff --git a/tests/asyncapi/kafka/test_app.py b/tests/asyncapi/kafka/test_app.py
index 4a79ad1e..5417db30 100644
--- a/tests/asyncapi/kafka/test_app.py
+++ b/tests/asyncapi/kafka/test_app.py
@@ -1,3 +1,5 @@
+from dirty_equals import IsPartialDict
+
 from faststream import FastStream
 from faststream.asyncapi.generate import get_app_schema
 from faststream.asyncapi.schema import Contact, ExternalDocs, License, Tag
@@ -148,11 +150,23 @@ def test_extra():
             title="My App",
             version="1.0.0",
             description="Test description",
-            license={"name": "MIT", "url": "https://mit.com/", "x-field": "extra"},
+            license={
+                "name": "MIT",
+                "url": "https://mit.com/",
+                "x-field": "extra",
+            },
             terms_of_service="https://my-terms.com/",
-            contact={"name": "support", "url": "https://help.com/", "x-field": "extra"},
+            contact={
+                "name": "support",
+                "url": "https://help.com/",
+                "x-field": "extra",
+            },
             tags=(
-                {"name": "some-tag", "description": "experimental", "x-field": "extra"},
+                {
+                    "name": "some-tag",
+                    "description": "experimental",
+                    "x-field": "extra",
+                },
             ),
             identifier="some-unique-uuid",
             external_docs={
@@ -162,33 +176,42 @@ def test_extra():
         )
     ).to_jsonable()
 
-    assert schema == {
-        "asyncapi": "2.6.0",
-        "channels": {},
-        "components": {"messages": {}, "schemas": {}},
-        "defaultContentType": "application/json",
-        "externalDocs": {"url": "https://extra-docs.py/", "x-field": "extra"},
-        "id": "some-unique-uuid",
-        "info": {
-            "contact": {
-                "name": "support",
-                "url": "https://help.com/",
+    assert schema == IsPartialDict(
+        {
+            "externalDocs": {
+                "url": "https://extra-docs.py/",
                 "x-field": "extra",
             },
-            "description": "Test description",
-            "license": {"name": "MIT", "url": "https://mit.com/", "x-field": "extra"},
-            "termsOfService": "https://my-terms.com/",
-            "title": "My App",
-            "version": "1.0.0",
-        },
-        "servers": {
-            "development": {
-                "protocol": "kafka",
-                "protocolVersion": "auto",
-                "url": "localhost",
-            }
-        },
-        "tags": [
-            {"description": "experimental", "name": "some-tag", "x-field": "extra"}
-        ],
-    }
+            "id": "some-unique-uuid",
+            "info": {
+                "contact": {
+                    "name": "support",
+                    "url": "https://help.com/",
+                    "x-field": "extra",
+                },
+                "description": "Test description",
+                "license": {
+                    "name": "MIT",
+                    "url": "https://mit.com/",
+                    "x-field": "extra",
+                },
+                "termsOfService": "https://my-terms.com/",
+                "title": "My App",
+                "version": "1.0.0",
+            },
+            "servers": {
+                "development": {
+                    "protocol": "kafka",
+                    "protocolVersion": "auto",
+                    "url": "localhost",
+                }
+            },
+            "tags": [
+                {
+                    "description": "experimental",
+                    "name": "some-tag",
+                    "x-field": "extra",
+                }
+            ],
+        }
+    )