Support fragments on unions.

tests/documents/unions/test.graphql

@@ -1,10 +1,55 @@
 query Nana {
   hallo {
+    __typename
     ... on Bar {
       nana
     }
     ... on Foo {
       forward
+      blip
     }
   }
 }
+
+query Nana2 {
+  hallo {
+    ...BazFragment
+  }
+}
+
+query Nana3 {
+  hallo {
+    ...VeryNestedFragment
+  }
+}
+
+query Nana4 {
+  hallo {
+    ...DelegatingFragment
+  }
+}
+
+fragment BazFragment on Element {
+  ... on Baz {
+    __typename
+    bloop
+  }
+}
+
+fragment VeryNestedFragment on Element {
+  __typename
+  ...BazFragment
+  ... on Bar {
+    nana
+  }
+}
+
+fragment DelegateFragment on Bar {
+  nana
+}
+
+fragment DelegatingFragment on Element {
+  ... on Bar {
+    ...DelegateFragment
+  }
+}

tests/schemas/union.graphql

@@ -1,15 +1,31 @@
+enum TestEnum1 {
+  A
+  B
+}
+
+enum TestEnum2 {
+  C
+  D
+}
+
 "This is foo"
 type Foo {
   "This is a forward ref"
   forward: String!
+  blip: TestEnum1
 }
 
 type Bar {
   "This is a forward ref"
   nana: Int!
 }
 
-union Element = Foo | Bar
+type Baz {
+  bloop: TestEnum2!
+}
+
+"This is a union"
+union Element = Foo | Bar | Baz
 
 type Query {
   hallo: Element

tests/test_referencer.py

@@ -1,5 +1,6 @@
 from turms.referencer import create_reference_registry_from_documents
 from turms.utils import parse_documents
+
 from .utils import build_relative_glob
 
 
@@ -19,3 +20,12 @@ def test_referencer_countries(countries_schema):
     assert (
         "StringQueryOperatorInput" not in z.inputs
     ), "StringQueryOperatorInput should be skipped"
+
+
+def test_referencer_enum_in_union_fragment(union_schema):
+
+    x = build_relative_glob("/documents/unions/*.graphql")
+    docs = parse_documents(union_schema, x)
+    z = create_reference_registry_from_documents(union_schema, docs)
+    assert "TestEnum1" in z.enums, "TestEnum1 should be referenced (in operation)"
+    assert "TestEnum2" in z.enums, "TestEnum2 should be referenced (in fragment)"

tests/test_unions.py

@@ -1,19 +1,19 @@
 
-from .utils import build_relative_glob, unit_test_with
 from turms.config import GeneratorConfig
-from turms.run import generate_ast
 from turms.plugins.enums import EnumsPlugin
-from turms.plugins.inputs import InputsPlugin
 from turms.plugins.fragments import FragmentsPlugin
-from turms.plugins.operations import OperationsPlugin
 from turms.plugins.funcs import (
-    FunctionDefinition,
     FuncsPlugin,
     FuncsPluginConfig,
+    FunctionDefinition,
 )
-from turms.stylers.snake_case import SnakeCaseStyler
-from turms.stylers.capitalize import CapitalizeStyler
+from turms.plugins.inputs import InputsPlugin
+from turms.plugins.operations import OperationsPlugin
 from turms.run import generate_ast
+from turms.stylers.capitalize import CapitalizeStyler
+from turms.stylers.snake_case import SnakeCaseStyler
+
+from .utils import build_relative_glob, unit_test_with
 
 
 def test_nested_input_funcs(union_schema):
@@ -52,3 +52,29 @@ def test_nested_input_funcs(union_schema):
         generated_ast,
         'Nana(hallo={"__typename": "Foo","forward": "yes"}).hallo.forward',
     )
+
+def test_fields_in_inline_fragment_on_union(union_schema):
+    config = GeneratorConfig(
+        documents=build_relative_glob("/documents/unions/*.graphql"),
+    )
+    generated_ast = generate_ast(
+        config,
+        union_schema,
+        stylers=[CapitalizeStyler(), SnakeCaseStyler()],
+        plugins=[EnumsPlugin(), InputsPlugin(), FragmentsPlugin(), OperationsPlugin()],
+    )
+
+    unit_test_with(
+        generated_ast,
+        """
+        assert Nana(hallo={'__typename': 'Foo', 'blip': 'A', 'forward': 'yes'}).hallo.blip == 'A'
+        assert Nana(hallo={'__typename': 'Bar', 'nana': 1}).hallo.nana == 1
+        assert Nana2(hallo={'__typename': 'Baz', 'bloop': 'C'}).hallo.bloop == 'C'
+        assert Nana3(hallo={'__typename': 'Baz', 'bloop': 'C'}).hallo.bloop == 'C'
+        assert Nana3(hallo={'__typename': 'Bar', 'nana': 1}).hallo.nana == 1
+        assert Nana4(hallo={'__typename': 'Bar', 'nana': 1}).hallo.nana == 1
+        """,
+    )
+
+
+

tests/utils.py

@@ -1,10 +1,12 @@
-import subprocess
-import os
 import ast
+import os
+import subprocess
 import sys
+import tempfile
+from textwrap import dedent
 from typing import List
+
 from turms.run import write_code_to_file
-import tempfile
 
 DIR_NAME = os.path.dirname(os.path.realpath(__file__))
 
@@ -65,7 +67,7 @@ def parse_to_code(tree: List[ast.AST]) -> str:
 
 def unit_test_with(generated_ast: List[ast.AST], test_string: str):
 
-    added_code = ast.parse(test_string).body
+    added_code = ast.parse(dedent(test_string)).body
     # We need to unparse before otherwise there might be complaints with missing lineno
     parsed_code = parse_to_code(generated_ast + added_code)