test: add test for discriminated unions

Signed-off-by: Luka Peschke <[email protected]>

Author: lukapeschke

tests/serializers/test_union.py

@@ -790,51 +790,92 @@ def __init__(self, type_: Literal['cat']) -> None:
         self.type_ = 'cat'
 
 
-def test_union_of_unions_of_models() -> None:
+@pytest.fixture
+def model_a_b_union_schema() -> core_schema.UnionSchema:
+    return core_schema.union_schema(
+        [
+            core_schema.model_schema(
+                cls=ModelA,
+                schema=core_schema.model_fields_schema(
+                    fields={
+                        'a': core_schema.model_field(core_schema.str_schema()),
+                        'b': core_schema.model_field(core_schema.str_schema()),
+                    },
+                ),
+            ),
+            core_schema.model_schema(
+                cls=ModelB,
+                schema=core_schema.model_fields_schema(
+                    fields={
+                        'c': core_schema.model_field(core_schema.str_schema()),
+                        'd': core_schema.model_field(core_schema.str_schema()),
+                    },
+                ),
+            ),
+        ]
+    )
+
+
+def test_union_of_unions_of_models(model_a_b_union_schema: core_schema.UnionSchema) -> None:
     s = SchemaSerializer(
         core_schema.union_schema(
             [
+                model_a_b_union_schema,
                 core_schema.union_schema(
                     [
                         core_schema.model_schema(
-                            cls=ModelA,
+                            cls=ModelCat,
                             schema=core_schema.model_fields_schema(
                                 fields={
-                                    'a': core_schema.model_field(core_schema.str_schema()),
-                                    'b': core_schema.model_field(core_schema.str_schema()),
+                                    'type_': core_schema.model_field(core_schema.literal_schema(['cat'])),
                                 },
                             ),
                         ),
                         core_schema.model_schema(
-                            cls=ModelB,
+                            cls=ModelDog,
                             schema=core_schema.model_fields_schema(
                                 fields={
-                                    'c': core_schema.model_field(core_schema.str_schema()),
-                                    'd': core_schema.model_field(core_schema.str_schema()),
+                                    'type_': core_schema.model_field(core_schema.literal_schema(['dog'])),
                                 },
                             ),
                         ),
                     ]
                 ),
-                core_schema.union_schema(
-                    [
-                        core_schema.model_schema(
+            ]
+        )
+    )
+
+    assert s.to_python(ModelA(a='a', b='b'), warnings='error') == {'a': 'a', 'b': 'b'}
+    assert s.to_python(ModelB(c='c', d='d'), warnings='error') == {'c': 'c', 'd': 'd'}
+    assert s.to_python(ModelCat(type_='cat'), warnings='error') == {'type_': 'cat'}
+    assert s.to_python(ModelDog(type_='dog'), warnings='error') == {'type_': 'dog'}
+
+
+def test_union_of_unions_of_models_with_tagged_union(model_a_b_union_schema: core_schema.UnionSchema) -> None:
+    s = SchemaSerializer(
+        core_schema.union_schema(
+            [
+                model_a_b_union_schema,
+                core_schema.tagged_union_schema(
+                    discriminator='type_',
+                    choices={
+                        'cat': core_schema.model_schema(
                             cls=ModelCat,
                             schema=core_schema.model_fields_schema(
                                 fields={
                                     'type_': core_schema.model_field(core_schema.literal_schema(['cat'])),
                                 },
                             ),
                         ),
-                        core_schema.model_schema(
+                        'dog': core_schema.model_schema(
                             cls=ModelDog,
                             schema=core_schema.model_fields_schema(
                                 fields={
                                     'type_': core_schema.model_field(core_schema.literal_schema(['dog'])),
                                 },
                             ),
                         ),
-                    ]
+                    },
                 ),
             ]
         )