Support for enum in unions (#152)

* feat: Enums in union implements interface

* chore: unused imports

Author: danilo-barboza

src/main/java/com/kobylynskyi/graphql/codegen/mapper/EnumDefinitionToDataModelMapper.java

@@ -3,14 +3,24 @@
 import com.kobylynskyi.graphql.codegen.model.EnumValueDefinition;
 import com.kobylynskyi.graphql.codegen.model.MappingContext;
 import com.kobylynskyi.graphql.codegen.model.definitions.ExtendedEnumTypeDefinition;
+import com.kobylynskyi.graphql.codegen.model.definitions.ExtendedUnionTypeDefinition;
 import graphql.language.Comment;
 import graphql.language.Directive;
 import graphql.language.DirectivesContainer;
 
-import java.util.*;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Collectors;
 
-import static com.kobylynskyi.graphql.codegen.model.DataModelFields.*;
+import static com.kobylynskyi.graphql.codegen.model.DataModelFields.CLASS_NAME;
+import static com.kobylynskyi.graphql.codegen.model.DataModelFields.FIELDS;
+import static com.kobylynskyi.graphql.codegen.model.DataModelFields.IMPLEMENTS;
+import static com.kobylynskyi.graphql.codegen.model.DataModelFields.JAVA_DOC;
+import static com.kobylynskyi.graphql.codegen.model.DataModelFields.PACKAGE;
 
 /**
  * Map enum definition to a Freemarker data model
@@ -31,11 +41,22 @@ public static Map<String, Object> map(MappingContext mappingContext, ExtendedEnu
         // type/enum/input/interface/union classes do not require any imports
         dataModel.put(PACKAGE, MapperUtils.getModelPackageName(mappingContext));
         dataModel.put(CLASS_NAME, MapperUtils.getClassNameWithPrefixAndSuffix(mappingContext, definition));
+        dataModel.put(IMPLEMENTS, getUnionInterfaces(mappingContext, definition));
         dataModel.put(JAVA_DOC, definition.getJavaDoc());
         dataModel.put(FIELDS, map(definition.getValueDefinitions()));
         return dataModel;
     }
 
+    private static Set<String> getUnionInterfaces(MappingContext mappingContext,
+                                                  ExtendedEnumTypeDefinition definition) {
+        return mappingContext.getDocument().getUnionDefinitions()
+                .stream()
+                .filter(union -> union.isDefinitionPartOfUnion(definition))
+                .map(ExtendedUnionTypeDefinition::getName)
+                .map(unionName -> MapperUtils.getClassNameWithPrefixAndSuffix(mappingContext, unionName))
+                .collect(Collectors.toSet());
+    }
+
     /**
      * Mapper from GraphQL's EnumValueDefinition to a Freemarker-understandable format
      *

src/main/resources/templates/javaClassGraphqlEnum.ftl

@@ -9,7 +9,7 @@ package ${package};
 </#list>
  */
 </#if>
-public enum ${className} {
+public enum ${className}<#if implements?has_content> implements <#list implements as interface>${interface}<#if interface_has_next>, </#if></#list></#if> {
 
 <#list fields as field>
 <#if field.javaDoc?has_content>

src/test/java/com/kobylynskyi/graphql/codegen/GraphQLCodegenUnionWithEnumTest.java

@@ -0,0 +1,55 @@
+package com.kobylynskyi.graphql.codegen;
+
+import com.kobylynskyi.graphql.codegen.model.MappingConfig;
+import com.kobylynskyi.graphql.codegen.utils.Utils;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+
+import static com.kobylynskyi.graphql.codegen.TestUtils.assertSameTrimmedContent;
+import static java.util.stream.Collectors.toList;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class GraphQLCodegenUnionWithEnumTest {
+
+    private final File outputBuildDir = new File("build/generated");
+    private final File outputJavaClassesDir = new File("build/generated/com/kobylynskyi/graphql/enumunion");
+
+    private GraphQLCodegen generator;
+
+    @BeforeEach
+    void init() {
+        MappingConfig mappingConfig = new MappingConfig();
+        mappingConfig.setPackageName("com.kobylynskyi.graphql.enumunion");
+        List<String> schemas = Collections.singletonList(
+                "src/test/resources/schemas/union-with-enum.graphqls"
+        );
+        generator = new GraphQLCodegen(schemas, outputBuildDir, mappingConfig);
+    }
+
+    @AfterEach
+    void cleanup() {
+        Utils.deleteDir(new File("build/generated"));
+    }
+
+    @Test
+    void generate_CheckFiles() throws Exception {
+        generator.generate();
+
+        File[] files = Objects.requireNonNull(outputJavaClassesDir.listFiles());
+        List<String> generatedFileNames = Arrays.stream(files).map(File::getName).sorted().collect(toList());
+        assertEquals(Arrays.asList("EnumMember1.java", "EnumMember2.java", "EnumUnion.java"), generatedFileNames);
+
+        for (File file : files) {
+            assertSameTrimmedContent(
+                    new File(String.format("src/test/resources/expected-classes/enum-union/%s.txt", file.getName())),
+                    file);
+        }
+    }
+}

src/test/resources/expected-classes/enum-union/EnumMember1.java.txt

@@ -0,0 +1,7 @@
+package com.kobylynskyi.graphql.enumunion;
+
+public enum EnumMember1 implements EnumUnion {
+
+    VALUE
+
+}

src/test/resources/expected-classes/enum-union/EnumMember2.java.txt

@@ -0,0 +1,7 @@
+package com.kobylynskyi.graphql.enumunion;
+
+public enum EnumMember2 implements EnumUnion {
+
+    OTHER_VALUE
+
+}

src/test/resources/expected-classes/enum-union/EnumUnion.java.txt

@@ -0,0 +1,6 @@
+package com.kobylynskyi.graphql.enumunion;
+
+
+public interface EnumUnion {
+
+}

src/test/resources/schemas/union-with-enum.graphqls

@@ -0,0 +1,9 @@
+enum EnumMember1 {
+    VALUE
+}
+
+enum EnumMember2 {
+    OTHER_VALUE
+}
+
+union EnumUnion = EnumMember1 | EnumMember2