functional tests for union type fix

Author: eli-bl

end_to_end_tests/functional_tests/generated_code_execution/test_enums_and_consts.py

@@ -133,88 +133,22 @@ def test_invalid_values(self, MyModel):
 """
 components:
   schemas:
-    MyEnum:
-      type: string
-      enum: ["a", "b"]
-    MyEnumIncludingNull:
-      type: ["string", "null"]
-      enum: ["a", "b", null]
-    MyNullOnlyEnum:
+    EnumOfNullOnly:
       enum: [null]
     MyModel:
       properties:
-        nullableEnumProp:
-          oneOf:
-            - {"$ref": "#/components/schemas/MyEnum"}
-            - type: "null"
-        enumIncludingNullProp: {"$ref": "#/components/schemas/MyEnumIncludingNull"}
-        nullOnlyEnumProp: {"$ref": "#/components/schemas/MyNullOnlyEnum"}
+        nullOnlyEnumProp: {"$ref": "#/components/schemas/EnumOfNullOnly"}
+      required: ["nullOnlyEnumProp"]
 """)
 @with_generated_code_imports(
-    ".models.MyEnum",
-    ".models.MyEnumIncludingNull",
     ".models.MyModel",
-    ".types.Unset",
 )
-class TestNullableEnums:
-    def test_nullable_enum_prop(self, MyModel, MyEnum, MyEnumIncludingNull):
-        assert_model_decode_encode(MyModel, {"nullableEnumProp": "b"}, MyModel(nullable_enum_prop=MyEnum.B))
-        assert_model_decode_encode(MyModel, {"nullableEnumProp": None}, MyModel(nullable_enum_prop=None))
-        assert_model_decode_encode(
-            MyModel,
-            {"enumIncludingNullProp": "a"},
-            MyModel(enum_including_null_prop=MyEnumIncludingNull.A),
-        )
-        assert_model_decode_encode( MyModel, {"enumIncludingNullProp": None}, MyModel(enum_including_null_prop=None))
+class TestSingleValueNullEnum:
+    def test_enum_of_null_only(self, MyModel):
         assert_model_decode_encode(MyModel, {"nullOnlyEnumProp": None}, MyModel(null_only_enum_prop=None))
 
-    def test_type_hints(self, MyModel, MyEnum, MyEnumIncludingNull, Unset):
-        assert_model_property_type_hint(MyModel, "nullable_enum_prop", Union[MyEnum, None, Unset])
-        assert_model_property_type_hint(MyModel, "enum_including_null_prop", Union[MyEnumIncludingNull, None, Unset])
-        assert_model_property_type_hint(MyModel, "null_only_enum_prop", Union[None, Unset])
-
-
-@with_generated_client_fixture(
-"""
-openapi: 3.0.0
-
-components:
-  schemas:
-    MyEnum:
-      type: string
-      enum: ["a", "b"]
-    MyEnumIncludingNull:
-      type: string
-      nullable: true
-      enum: ["a", "b", null]
-    MyModel:
-      properties:
-        nullableEnumProp:
-          allOf:
-            - {"$ref": "#/components/schemas/MyEnum"}
-          nullable: true
-        enumIncludingNullProp: {"$ref": "#/components/schemas/MyEnumIncludingNull"}
-""")
-@with_generated_code_imports(
-    ".models.MyEnum",
-    ".models.MyEnumIncludingNull",
-    ".models.MyModel",
-    ".types.Unset",
-)
-class TestNullableEnumsInOpenAPI30:
-    def test_nullable_enum_prop(self, MyModel, MyEnum, MyEnumIncludingNull):
-        assert_model_decode_encode(MyModel, {"nullableEnumProp": "b"}, MyModel(nullable_enum_prop=MyEnum.B))
-        assert_model_decode_encode(MyModel, {"nullableEnumProp": None}, MyModel(nullable_enum_prop=None))
-        assert_model_decode_encode(
-            MyModel,
-            {"enumIncludingNullProp": "a"},
-            MyModel(enum_including_null_prop=MyEnumIncludingNull.A),
-        )
-        assert_model_decode_encode( MyModel, {"enumIncludingNullProp": None}, MyModel(enum_including_null_prop=None))
-    
-    def test_type_hints(self, MyModel, MyEnum, MyEnumIncludingNull, Unset):
-        assert_model_property_type_hint(MyModel, "nullable_enum_prop", Union[MyEnum, None, Unset])
-        assert_model_property_type_hint(MyModel, "enum_including_null_prop", Union[MyEnumIncludingNull, None, Unset])
+    def test_type_hints(self, MyModel):
+        assert_model_property_type_hint(MyModel, "null_only_enum_prop", None)
 
 
 @with_generated_client_fixture(
@@ -259,6 +193,8 @@ def test_invalid_int(self, MyModel):
 
 @with_generated_client_fixture(
 """
+# Tests of literal_enums mode, where enums become a typing.Literal type instead of a class
+
 components:
   schemas:
     MyEnum:
@@ -303,6 +239,8 @@ def test_invalid_values(self, MyModel):
 
 @with_generated_client_fixture(
 """
+# Tests of literal_enums mode, where enums become a typing.Literal type instead of a class
+
 components:
   schemas:
     MyEnum:
@@ -347,6 +285,8 @@ def test_invalid_values(self, MyModel):
 
 @with_generated_client_fixture(
 """
+# Similar to some of the "union with null" tests in test_unions.py, but in literal_enums mode
+
 components:
   schemas:
     MyEnum:

end_to_end_tests/functional_tests/generated_code_execution/test_unions.py

@@ -38,6 +38,8 @@ def test_type_hints(self, MyModel, Unset):
 
 @with_generated_client_fixture(
 """
+# Various use cases for oneOf
+
 components:
   schemas:
     ThingA:
@@ -152,6 +154,203 @@ def test_type_hints(self, ModelWithUnion, ModelWithRequiredUnion, ModelWithUnion
         )
 
 
+@with_generated_client_fixture(
+"""
+# Various use cases for a oneOf where one of the variants is null, since these are handled
+# a bit differently in the generator
+
+components:
+  schemas:
+    MyEnum:
+      type: string
+      enum: ["a", "b"]
+    MyObject:
+      type: object
+      properties:
+        name:
+          type: string
+    MyModel:
+      properties:
+        nullableEnumProp:
+          oneOf:
+            - {"$ref": "#/components/schemas/MyEnum"}
+            - type: "null"
+        nullableObjectProp:
+          oneOf:
+            - {"$ref": "#/components/schemas/MyObject"}
+            - type: "null"
+        inlineNullableObject:
+          # Note, the generated class for this should be called "MyModelInlineNullableObject",
+          # since the generator's rule for inline schemas that require their own class is to
+          # concatenate the property name to the parent schema name.
+          oneOf:
+            - type: object
+              properties:
+                name:
+                  type: string
+            - type: "null"
+""")
+@with_generated_code_imports(
+    ".models.MyEnum",
+    ".models.MyObject",
+    ".models.MyModel",
+    ".models.MyModelInlineNullableObject",
+    ".types.Unset",
+)
+class TestUnionsWithNull:
+    def test_nullable_enum_prop(self, MyModel, MyEnum):
+        assert_model_decode_encode(MyModel, {"nullableEnumProp": "b"}, MyModel(nullable_enum_prop=MyEnum.B))
+        assert_model_decode_encode(MyModel, {"nullableEnumProp": None}, MyModel(nullable_enum_prop=None))
+
+    def test_nullable_object_prop(self, MyModel, MyObject):
+        assert_model_decode_encode( MyModel, {"nullableObjectProp": None}, MyModel(nullable_object_prop=None))
+        assert_model_decode_encode( MyModel, {"nullableObjectProp": None}, MyModel(nullable_object_prop=None))
+
+    def test_nullable_object_prop_with_inline_schema(self, MyModel, MyModelInlineNullableObject):
+        assert_model_decode_encode(
+            MyModel,
+            {"inlineNullableObject": {"name": "a"}},
+            MyModel(inline_nullable_object=MyModelInlineNullableObject(name="a")),
+        )
+        assert_model_decode_encode( MyModel, {"inlineNullableObject": None}, MyModel(inline_nullable_object=None))
+    
+    def test_type_hints(self, MyModel, MyEnum, Unset):
+        assert_model_property_type_hint(MyModel, "nullable_enum_prop", Union[MyEnum, None, Unset])
+        assert_model_property_type_hint(MyModel, "nullable_object_prop", Union[ForwardRef("MyObject"), None, Unset])
+        assert_model_property_type_hint(
+            MyModel,
+            "inline_nullable_object",
+            Union[ForwardRef("MyModelInlineNullableObject"), None, Unset],
+        )
+
+
+@with_generated_client_fixture(
+"""
+# Tests for combining the OpenAPI 3.0 "nullable" attribute with an enum
+
+openapi: 3.0.0
+
+components:
+  schemas:
+    MyEnum:
+      type: string
+      enum: ["a", "b"]
+    MyEnumIncludingNull:
+      type: string
+      nullable: true
+      enum: ["a", "b", null]
+    MyModel:
+      properties:
+        nullableEnumProp:
+          allOf:
+            - {"$ref": "#/components/schemas/MyEnum"}
+          nullable: true
+        enumIncludingNullProp: {"$ref": "#/components/schemas/MyEnumIncludingNull"}
+""")
+@with_generated_code_imports(
+    ".models.MyEnum",
+    ".models.MyEnumIncludingNull",
+    ".models.MyModel",
+    ".types.Unset",
+)
+class TestNullableEnumsInOpenAPI30:
+    def test_nullable_enum_prop(self, MyModel, MyEnum, MyEnumIncludingNull):
+        assert_model_decode_encode(MyModel, {"nullableEnumProp": "b"}, MyModel(nullable_enum_prop=MyEnum.B))
+        assert_model_decode_encode(MyModel, {"nullableEnumProp": None}, MyModel(nullable_enum_prop=None))
+        assert_model_decode_encode(
+            MyModel,
+            {"enumIncludingNullProp": "a"},
+            MyModel(enum_including_null_prop=MyEnumIncludingNull.A),
+        )
+        assert_model_decode_encode( MyModel, {"enumIncludingNullProp": None}, MyModel(enum_including_null_prop=None))
+    
+    def test_type_hints(self, MyModel, MyEnum, MyEnumIncludingNull, Unset):
+        assert_model_property_type_hint(MyModel, "nullable_enum_prop", Union[MyEnum, None, Unset])
+        assert_model_property_type_hint(MyModel, "enum_including_null_prop", Union[MyEnumIncludingNull, None, Unset])
+
+
+@with_generated_client_fixture(
+"""
+# Test use cases where there's a union of types *and* an explicit list of multiple "type:"s
+
+components:
+  schemas:
+    MyStringEnum:
+      type: string
+      enum: ["a", "b"]
+    MyIntEnum:
+      type: integer
+      enum: [1, 2]
+    MyEnumIncludingNull:
+      type: ["string", "null"]
+      enum: ["a", "b", null]
+    MyObject:
+      type: object
+      properties:
+        name:
+          type: string
+    MyModel:
+      properties:
+        enumsWithListOfTypesProp:
+          type: ["string", "integer"]
+          oneOf:
+            - {"$ref": "#/components/schemas/MyStringEnum"}
+            - {"$ref": "#/components/schemas/MyIntEnum"}
+        enumIncludingNullProp: {"$ref": "#/components/schemas/MyEnumIncludingNull"}
+        nullableObjectWithListOfTypesProp:
+          type: ["string", "object"]
+          oneOf:
+            - {"$ref": "#/components/schemas/MyObject"}
+            - type: "null"
+""")
+@with_generated_code_imports(
+    ".models.MyStringEnum",
+    ".models.MyIntEnum",
+    ".models.MyEnumIncludingNull",
+    ".models.MyObject",
+    ".models.MyModel",
+    ".types.Unset",
+)
+class TestUnionsWithExplicitListOfTypes:
+    # This covers some use cases where combining "oneOf" with "type: [list of types]" (which is fine
+    # to do in OpenAPI) used to generate enum/model classes incorrectly.
+
+    def test_union_of_enums(self, MyModel, MyStringEnum, MyIntEnum):
+        assert_model_decode_encode(
+            MyModel,
+            {"enumsWithListOfTypesProp": "b"},
+            MyModel(enums_with_list_of_types_prop=MyStringEnum.B),
+        )
+        assert_model_decode_encode(
+            MyModel,
+            {"enumsWithListOfTypesProp": 2},
+            MyModel(enums_with_list_of_types_prop=MyIntEnum.VALUE_2),
+        )
+
+    def test_union_of_enum_with_null(self, MyModel, MyEnumIncludingNull):
+        assert_model_decode_encode(
+            MyModel,
+            {"enumIncludingNullProp": "b"},
+            MyModel(enum_including_null_prop=MyEnumIncludingNull.B),
+        )
+        assert_model_decode_encode(
+            MyModel,
+            {"enumIncludingNullProp": None},
+            MyModel(enum_including_null_prop=None),
+        )
+
+    def test_nullable_object_with_list_of_types(self, MyModel, MyObject):
+        assert_model_decode_encode(
+            MyModel,
+            {"nullableObjectWithListOfTypesProp": {"name": "a"}},
+            MyModel(nullable_object_with_list_of_types_prop=MyObject(name="a")),
+        )
+        assert_model_decode_encode(
+            MyModel,
+            {"nullableObjectWithListOfTypesProp": None},
+            MyModel(nullable_object_with_list_of_types_prop=None),
+        )
+    
 @with_generated_client_fixture(
 """
 components:

end_to_end_tests/functional_tests/helpers.py

@@ -47,8 +47,7 @@ def _decorator(cls):
         nonlocal alias
 
         def _func(self, generated_client):
-            module = generated_client.import_module(module_name)
-            return getattr(module, import_name)
+            return generated_client.import_symbol(module_name, import_name)
 
         alias = alias or import_name
         _func.__name__ = alias

end_to_end_tests/generated_client.py

@@ -45,6 +45,16 @@ def import_module(self, module_path: str) -> Any:
         """Attempt to import a module from the generated code."""
         return importlib.import_module(f"{self.base_module}{module_path}")
 
+    def import_symbol(self, module_path: str, name: str) -> Any:
+        module = self.import_module(module_path)
+        try:
+            return getattr(module, name)
+        except AttributeError:
+            existing = ", ".join(name for name in dir(module) if not name.startswith("_"))
+            assert False, (
+                f"Couldn't find import \"{name}\" in \"{self.base_module}{module_path}\"."
+                f" Available imports in that module are: {existing}"
+            )
 
 def _run_command(
     command: str,