more discriminator tests

Author: eli-bl

end_to_end_tests/functional_tests/generated_code_execution/test_unions.py

@@ -167,6 +167,20 @@ def test_type_hints(self, ModelWithUnion, ModelWithRequiredUnion, ModelWithUnion
         modelType: {"type": "string"}
         name: {"type": "string"}
       required: ["modelType"]
+    Corgi:
+      type: object
+      properties:
+        modelType: {"type": "string"}
+        dogType: {"type": "string"}
+        name: {"type": "string"}
+      required: ["modelType"]
+    Schnauzer:
+      type: object
+      properties:
+        modelType: {"type": "string"}
+        dogType: {"type": "string"}
+        name: {"type": "string"}
+      required: ["modelType"]
     WithDiscriminatedUnion:
       type: object
       properties:
@@ -178,8 +192,8 @@ def test_type_hints(self, ModelWithUnion, ModelWithRequiredUnion, ModelWithUnion
             propertyName: modelType
             mapping:
               "type1": "#/components/schemas/ModelType1"
-              "type2": "#/components/schemas/ModelType2"
-              "another-value": "#/components/schemas/ModelType2"
+              "type2": "ModelType2"  # exactly equivalent to "#/components/schemas/ModelType2"
+              "another-value": "#/components/schemas/ModelType2"  # deliberately mapped to same type as previous line
     WithDiscriminatedUnionImplicitMapping:
       type: object
       properties:
@@ -189,12 +203,46 @@ def test_type_hints(self, ModelWithUnion, ModelWithRequiredUnion, ModelWithUnion
             - $ref: "#/components/schemas/ModelType2"
           discriminator:
             propertyName: modelType
+    WithNestedDiscriminatorsSameProperty:
+      type: object
+      properties:
+        unionProp:
+          oneOf:
+            - oneOf:
+              - $ref: "#/components/schemas/ModelType1"
+              - $ref: "#/components/schemas/ModelType2"
+              discriminator:
+                propertyName: modelType
+            - oneOf:
+              - $ref: "#/components/schemas/Corgi"
+              - $ref: "#/components/schemas/Schnauzer"
+              discriminator:
+                propertyName: modelType
+    WithNestedDiscriminatorsDifferentProperty:
+      type: object
+      properties:
+        unionProp:
+          oneOf:
+            - oneOf:
+              - $ref: "#/components/schemas/ModelType1"
+              - $ref: "#/components/schemas/ModelType2"
+              discriminator:
+                propertyName: modelType
+            - oneOf:
+              - $ref: "#/components/schemas/Corgi"
+              - $ref: "#/components/schemas/Schnauzer"
+              discriminator:
+                propertyName: dogType
 """)
 @with_generated_code_imports(
     ".models.ModelType1",
     ".models.ModelType2",
+    ".models.Corgi",
+    ".models.Schnauzer",
     ".models.WithDiscriminatedUnion",
     ".models.WithDiscriminatedUnionImplicitMapping",
+    ".models.WithNestedDiscriminatorsSameProperty",
+    ".models.WithNestedDiscriminatorsDifferentProperty",
 )
 class TestDiscriminators:
     def test_with_explicit_mapping(self, ModelType1, ModelType2, WithDiscriminatedUnion):
@@ -231,3 +279,27 @@ def test_with_implicit_mapping(self, ModelType1, ModelType2, WithDiscriminatedUn
         )
         with pytest.raises(TypeError):
             WithDiscriminatedUnionImplicitMapping.from_dict({"unionProp": {"modelType": "unknown-value"}})
+
+    def test_nested_with_same_property(self, ModelType1, Schnauzer, WithNestedDiscriminatorsSameProperty):
+        assert_model_decode_encode(
+            WithNestedDiscriminatorsSameProperty,
+            {"unionProp": {"modelType": "ModelType1", "name": "a"}},
+            WithNestedDiscriminatorsSameProperty(union_prop=ModelType1(model_type="ModelType1", name="a")),
+        )
+        assert_model_decode_encode(
+            WithNestedDiscriminatorsSameProperty,
+            {"unionProp": {"modelType": "Schnauzer", "name": "a"}},
+            WithNestedDiscriminatorsSameProperty(union_prop=Schnauzer(model_type="Schnauzer", name="a")),
+        )
+
+    def test_nested_with_different_property(self, ModelType1, Schnauzer, WithNestedDiscriminatorsDifferentProperty):
+        assert_model_decode_encode(
+            WithNestedDiscriminatorsDifferentProperty,
+            {"unionProp": {"modelType": "ModelType1", "name": "a"}},
+            WithNestedDiscriminatorsDifferentProperty(union_prop=ModelType1(model_type="ModelType1", name="a")),
+        )
+        assert_model_decode_encode(
+            WithNestedDiscriminatorsDifferentProperty,
+            {"unionProp": {"modelType": "irrelevant", "dogType": "Schnauzer", "name": "a"}},
+            WithNestedDiscriminatorsDifferentProperty(union_prop=Schnauzer(model_type="irrelevant", dog_type="Schnauzer", name="a")),
+        )

end_to_end_tests/functional_tests/generator_failure_cases/test_invalid_unions.py

@@ -32,4 +32,110 @@ def test_invalid_default(self, warnings):
         assert_bad_schema_warning(warnings, "UnionWithInvalidDefault", "Invalid int value: aaa")
 
     def test_invalid_property(self, warnings):
-        assert_bad_schema_warning(warnings, "UnionWithMalformedVariant", "Invalid property in union")
+        assert_bad_schema_warning(warnings, "UnionWithMalformedVariant", "Invalid property in union")
+
+
+class TestInvalidDiscriminators:
+    @pytest.fixture(scope="class")
+    def warnings(self):
+        return inline_spec_should_cause_warnings(
+"""
+components:
+  schemas:
+    ModelType1:
+      type: object
+      properties:
+        modelType: {"type": "string"}
+        name: {"type": "string"}
+      required: ["modelType"]
+    ModelType2:
+      type: object
+      properties:
+        modelType: {"type": "string"}
+        name: {"type": "string"}
+      required: ["modelType"]
+    ModelType3:
+      type: object
+      properties:
+        modelType: {"type": "string"}
+        name: {"type": "string"}
+      required: ["modelType"]
+    StringType:
+      type: string
+    WithUnknownSchemaInMapping:
+      type: object
+      properties:
+        unionProp:
+          oneOf:
+            - $ref: "#/components/schemas/ModelType1"
+            - $ref: "#/components/schemas/ModelType2"
+          discriminator:
+            propertyName: modelType
+            mapping:
+              "type1": "#/components/schemas/ModelType1"
+              "type2": "#/components/schemas/DoesntExist"
+    WithReferenceToSchemaNotInUnion:
+      type: object
+      properties:
+        unionProp:
+          oneOf:
+            - $ref: "#/components/schemas/ModelType1"
+            - $ref: "#/components/schemas/ModelType2"
+          discriminator:
+            propertyName: modelType
+            mapping:
+              "type1": "#/components/schemas/ModelType1"
+              "type2": "#/components/schemas/ModelType2"
+              "type3": "#/components/schemas/ModelType3"
+    WithNonObjectVariant:
+      type: object
+      properties:
+        unionProp:
+          oneOf:
+            - $ref: "#/components/schemas/ModelType1"
+            - $ref: "#/components/schemas/StringType"
+          discriminator:
+            propertyName: modelType
+    WithInlineSchema:
+      type: object
+      properties:
+        unionProp:
+          oneOf:
+            - $ref: "#/components/schemas/ModelType1"
+            - type: object
+              properties:
+                modelType: {"type": "string"}
+                name: {"type": "string"}
+          discriminator:
+            propertyName: modelType
+      
+"""
+        )
+    
+    def test_invalid_reference(self, warnings):
+        assert_bad_schema_warning(
+            warnings,
+            "WithUnknownSchemaInMapping",
+            'Invalid reference "#/components/schemas/DoesntExist" in discriminator mapping',
+        )
+
+    def test_reference_to_schema_not_in_union(self, warnings):
+        assert_bad_schema_warning(
+            warnings,
+            "WithReferenceToSchemaNotInUnion",
+            'Discriminator mapping referred to "ModelType3" which is not one of the schema variants',
+        )
+
+    def test_non_object_variant(self, warnings):
+        assert_bad_schema_warning(
+            warnings,
+            "WithNonObjectVariant",
+            "All schema variants must be objects when using a discriminator",
+        )
+
+    def test_inline_schema(self, warnings):
+        assert_bad_schema_warning(
+            warnings,
+            "WithInlineSchema",
+            "Inline schema declarations are not allowed when using a discriminator",
+        )