add a test

graphene/tests/issues/test_356.py

@@ -1,24 +0,0 @@
-# https://github.com/graphql-python/graphene/issues/356
-
-import pytest
-import graphene
-from graphene import relay
-
-class SomeTypeOne(graphene.ObjectType):
-    pass
-
-class SomeTypeTwo(graphene.ObjectType):
-    pass
-
-class MyUnion(graphene.Union):
-    class Meta:
-        types = (SomeTypeOne, SomeTypeTwo)
-
-def test_issue():
-    with pytest.raises(Exception) as exc_info:
-        class Query(graphene.ObjectType):
-            things = relay.ConnectionField(MyUnion)
-
-        schema = graphene.Schema(query=Query)
-
-    assert str(exc_info.value) == 'IterableConnectionField type have to be a subclass of Connection. Received "MyUnion".'

graphene/types/tests/test_union.py

@@ -1,7 +1,9 @@
 import pytest
 
 from ..objecttype import ObjectType
+from ..schema import Schema
 from ..union import Union
+from graphene.relay.connection import ConnectionField
 
 
 class MyObjectType1(ObjectType):
@@ -41,3 +43,33 @@ def test_generate_union_with_no_types():
             pass
 
     assert str(exc_info.value) == 'Must provide types for Union MyUnion.'
+
+
+def test_union_as_connection():
+    class MyUnion(Union):
+        class Meta:
+            types = (MyObjectType1, MyObjectType2)
+
+    class Query(ObjectType):
+        objects = ConnectionField(MyUnion)
+
+        def resolve_objects(self, args, context, info):
+            return [MyObjectType1(), MyObjectType2()]
+
+    query = '''
+    query {
+        objects {
+            edges {
+                node {
+                    __typename
+                }
+            }
+        }
+    }
+    '''
+    schema = Schema(query=Query)
+    result = schema.execute(query)
+    assert not result.errors
+    assert len(result.data['objects']['edges']) == 2
+    assert result.data['objects']['edges'][0]['node']['__typename'] == 'MyObjectType1'
+    assert result.data['objects']['edges'][1]['node']['__typename'] == 'MyObjectType2'