do not assume discriminator is not a property (#1811)

* do not assume discriminator is not a property * chore: fix precommit * add test_pydantic_model_with_keyword_property * rename argument to hide copy-paste origin * chore: fix precommit --------- Co-authored-by: Nikita Pastukhov <[email protected]>
airtai · Sep 25, 2024 · f3cdffd · f3cdffd
1 parent af92337
commit f3cdffd
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 1 deletion.
diff --git a/faststream/asyncapi/generate.py b/faststream/asyncapi/generate.py
@@ -212,7 +212,10 @@ def _move_pydantic_refs(
             for i in range(len(data[k])):
                 data[k][i] = _move_pydantic_refs(item[i], key)
 
-    if isinstance(desciminator := data.get("discriminator"), dict):
+    if (
+        isinstance(desciminator := data.get("discriminator"), dict)
+        and "propertyName" in desciminator
+    ):
         data["discriminator"] = desciminator["propertyName"]
 
     return data
diff --git a/tests/asyncapi/base/arguments.py b/tests/asyncapi/base/arguments.py
@@ -396,6 +396,33 @@ async def handle(user: User): ...
                 "type": "object",
             }
 
+    def test_pydantic_model_with_keyword_property(self):
+        class TestModel(pydantic.BaseModel):
+            discriminator: int = 0
+
+        broker = self.broker_class()
+
+        @broker.subscriber("test")
+        async def handle(model: TestModel): ...
+
+        schema = get_app_schema(self.build_app(broker)).to_jsonable()
+
+        payload = schema["components"]["schemas"]
+
+        for key, v in payload.items():
+            assert key == "TestModel"
+            assert v == {
+                "properties": {
+                    "discriminator": {
+                        "default": 0,
+                        "title": "Discriminator",
+                        "type": "integer",
+                    },
+                },
+                "title": key,
+                "type": "object",
+            }
+
     def test_ignores_depends(self):
         broker = self.broker_class()
 

diff --git a/tests/asyncapi/base/publisher.py b/tests/asyncapi/base/publisher.py
@@ -126,3 +126,29 @@ async def handler(msg: str):
         schema = get_app_schema(self.build_app(broker))
 
         assert schema.channels == {}, schema.channels
+
+    def test_pydantic_model_with_keyword_property_publisher(self):
+        class TestModel(pydantic.BaseModel):
+            discriminator: int = 0
+
+        broker = self.broker_class()
+
+        @broker.publisher("test")
+        async def handle(msg) -> TestModel: ...
+
+        schema = get_app_schema(self.build_app(broker)).to_jsonable()
+
+        payload = schema["components"]["schemas"]
+
+        for key, v in payload.items():
+            assert v == {
+                "properties": {
+                    "discriminator": {
+                        "default": 0,
+                        "title": "Discriminator",
+                        "type": "integer",
+                    },
+                },
+                "title": key,
+                "type": "object",
+            }