Improved Enum serialization tests

Author: syrusakbary

graphql/pyutils/compat.py

@@ -27,7 +27,7 @@
 try:
     from enum import Enum
 except ImportError:
-    from .enum import Enum
+    from .enum import Enum  # type: ignore
 
 if False:
     from typing import Callable

graphql/type/definition.py

@@ -516,7 +516,7 @@ def get_value(self, name):
         return self._name_lookup.get(name)
 
     def serialize(self, value):
-        # type: (str) -> Optional[str]
+        # type: (Union[str, PyEnum]) -> Optional[str]
         if isinstance(value, PyEnum):
             # We handle PyEnum values
             value = value.value

graphql/type/tests/test_serialization.py

@@ -1,6 +1,8 @@
 import pytest
 
-from graphql.type import GraphQLBoolean, GraphQLFloat, GraphQLInt, GraphQLString
+from ..scalars import GraphQLBoolean, GraphQLFloat, GraphQLInt, GraphQLString
+from ..definition import GraphQLEnumType, GraphQLEnumValue
+from ...pyutils.compat import Enum
 
 
 def test_serializes_output_int():
@@ -56,3 +58,26 @@ def test_serializes_output_boolean():
     assert GraphQLBoolean.serialize(0) is False
     assert GraphQLBoolean.serialize(True) is True
     assert GraphQLBoolean.serialize(False) is False
+
+
+def test_serializes_enum():
+    class Color(Enum):
+        RED = "RED"
+        GREEN = "GREEN"
+        BLUE = "BLUE"
+        EXTRA = "EXTRA"
+
+    enum_type = GraphQLEnumType(
+        "Color",
+        values={
+            "RED": GraphQLEnumValue("RED"),
+            "GREEN": GraphQLEnumValue("GREEN"),
+            "BLUE": GraphQLEnumValue("BLUE"),
+        },
+    )
+    assert enum_type.serialize("RED") == "RED"
+    assert enum_type.serialize("NON_EXISTING") is None
+    assert enum_type.serialize(Color.RED) == "RED"
+    assert enum_type.serialize(Color.RED.value) == "RED"
+    assert enum_type.serialize(Color.EXTRA) is None
+    assert enum_type.serialize(Color.EXTRA.value) is None