test: add basic example for generated schema

pydantic · Feb 9, 2021 · 78248a1 · 78248a1
1 parent 11b9cee
commit 78248a1
Showing 1 changed file with 65 additions and 0 deletions.
diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -2149,3 +2149,68 @@ class Model(BaseModel):
         f'{module_2.__name__}__MyEnum',
         f'{module_2.__name__}__MyModel',
     }
+
+
+@pytest.mark.skipif(not Literal, reason='typing_extensions not installed and python version < 3.8')
+def test_discriminated_union():
+    class Cat(BaseModel):
+        pet_type: Literal['cat']
+
+    class Dog(BaseModel):
+        pet_type: Literal['dog']
+
+    class Lizard(BaseModel):
+        pet_type: Literal['reptile', 'lizard']
+
+    class Model(BaseModel):
+        pet: Union[Cat, Dog, Lizard] = Field(..., discriminator='pet_type')
+
+    assert Model.schema() == {
+        'title': 'Model',
+        'type': 'object',
+        'properties': {
+            'pet': {
+                'title': 'Pet',
+                'anyOf': [
+                    {'$ref': '#/definitions/Cat'},
+                    {'$ref': '#/definitions/Dog'},
+                    {'$ref': '#/definitions/Lizard'},
+                ],
+                'discriminator': {
+                    'propertyName': 'pet_type',
+                    'mapping': {
+                        'cat': '#/definitions/Cat',
+                        'dog': '#/definitions/Dog',
+                        'lizard': '#/definitions/Lizard',
+                        'reptile': '#/definitions/Lizard',
+                    },
+                },
+            }
+        },
+        'required': ['pet'],
+        'definitions': {
+            'Cat': {
+                'title': 'Cat',
+                'type': 'object',
+                'properties': {'pet_type': {'const': 'cat', 'title': 'Pet Type', 'type': 'string'}},
+                'required': ['pet_type'],
+            },
+            'Dog': {
+                'title': 'Dog',
+                'type': 'object',
+                'properties': {'pet_type': {'const': 'dog', 'title': 'Pet Type', 'type': 'string'}},
+                'required': ['pet_type'],
+            },
+            'Lizard': {
+                'title': 'Lizard',
+                'type': 'object',
+                'properties': {
+                    'pet_type': {
+                        'anyOf': [{'const': 'reptile', 'type': 'string'}, {'const': 'lizard', 'type': 'string'}],
+                        'title': 'Pet Type',
+                    }
+                },
+                'required': ['pet_type'],
+            },
+        },
+    }