Add companion object generation for enum classes in kotlin codegen2 (N…

…etflix#670) * Add companion object generation to kotlin codegen2 * Fix indentation * Fix outdated "expected" values of tests
dwilkolek · Apr 15, 2024 · d33cbe2 · d33cbe2
1 parent 6c11193
commit d33cbe2
Show file tree

Hide file tree

Showing 8 changed files with 65 additions and 0 deletions.
diff --git a/...tegTest/kotlin/com/netflix/graphql/dgs/codegen/cases/enum/expected/types/EmployeeTypes.kt b/...tegTest/kotlin/com/netflix/graphql/dgs/codegen/cases/enum/expected/types/EmployeeTypes.kt
@@ -4,4 +4,7 @@ public enum class EmployeeTypes {
   ENGINEER,
   MANAGER,
   DIRECTOR,
+  ;
+
+  public companion object
 }
diff --git a/...c/integTest/kotlin/com/netflix/graphql/dgs/codegen/cases/enumDocs/expected/types/Color.kt b/...c/integTest/kotlin/com/netflix/graphql/dgs/codegen/cases/enumDocs/expected/types/Color.kt
@@ -7,4 +7,7 @@ public enum class Color {
   red,
   white,
   blue,
+  ;
+
+  public companion object
 }
diff --git a/...om/netflix/graphql/dgs/codegen/cases/enumWithExtendedType/expected/types/EmployeeTypes.kt b/...om/netflix/graphql/dgs/codegen/cases/enumWithExtendedType/expected/types/EmployeeTypes.kt
@@ -5,4 +5,7 @@ public enum class EmployeeTypes {
   MANAGER,
   DIRECTOR,
   QA,
+  ;
+
+  public companion object
 }
diff --git a/...tflix/graphql/dgs/codegen/cases/inputWithDefaultEnumValueForArray/expected/types/Color.kt b/...tflix/graphql/dgs/codegen/cases/inputWithDefaultEnumValueForArray/expected/types/Color.kt
@@ -3,4 +3,7 @@ package com.netflix.graphql.dgs.codegen.cases.inputWithDefaultEnumValueForArray.
 public enum class Color {
   red,
   blue,
+  ;
+
+  public companion object
 }
diff --git a/...om/netflix/graphql/dgs/codegen/cases/inputWithDefaultValueForEnum/expected/types/Color.kt b/...om/netflix/graphql/dgs/codegen/cases/inputWithDefaultValueForEnum/expected/types/Color.kt
@@ -2,4 +2,7 @@ package com.netflix.graphql.dgs.codegen.cases.inputWithDefaultValueForEnum.expec
 
 public enum class Color {
   red,
+  ;
+
+  public companion object
 }
diff --git a/...gTest/kotlin/com/netflix/graphql/dgs/codegen/cases/projectionWithEnum/expected/types/E.kt b/...gTest/kotlin/com/netflix/graphql/dgs/codegen/cases/projectionWithEnum/expected/types/E.kt
@@ -2,4 +2,7 @@ package com.netflix.graphql.dgs.codegen.cases.projectionWithEnum.expected.types
 
 public enum class E {
   V,
+  ;
+
+  public companion object
 }
diff --git a/...ain/kotlin/com/netflix/graphql/dgs/codegen/generators/kotlin2/GenerateKotlin2EnumTypes.kt b/...ain/kotlin/com/netflix/graphql/dgs/codegen/generators/kotlin2/GenerateKotlin2EnumTypes.kt
@@ -54,6 +54,10 @@ fun generateKotlin2EnumTypes(
                 .plus(extensionTypes)
                 .flatMap { it.enumValueDefinitions }
 
+            val companionObject = TypeSpec.companionObjectBuilder()
+                .addOptionalGeneratedAnnotation(config)
+                .build()
+
             // create the enum class
             val enumSpec = TypeSpec.classBuilder(enumDefinition.name)
                 .addOptionalGeneratedAnnotation(config)
@@ -80,6 +84,7 @@ fun generateKotlin2EnumTypes(
                             .build()
                     }
                 )
+                .addType(companionObject)
                 .build()
 
             // return a file per enum

diff --git a/...l-dgs-codegen-core/src/test/kotlin/com/netflix/graphql/dgs/codegen/Kotline2CodeGenTest.kt b/...l-dgs-codegen-core/src/test/kotlin/com/netflix/graphql/dgs/codegen/Kotline2CodeGenTest.kt
@@ -85,11 +85,53 @@ class Kotline2CodeGenTest {
                 |public enum class TownJobTypes {
                 |  @Deprecated(message = "town switched to electric lights")
                 |  LAMPLIGHTER,
+                |  ;
+                |
+                |  public companion object
                 |}
                 |
             """.trimMargin()
 
         )
         assertCompilesKotlin(result.kotlinEnumTypes)
     }
+
+    @Test
+    fun `Add companion object to enum class`() {
+        val schema = """
+            enum MyEnum {
+                A
+                B
+                C
+            }
+        """.trimIndent()
+
+        val result = CodeGen(
+            CodeGenConfig(
+                schemas = setOf(schema),
+                packageName = basePackageName,
+                language = Language.KOTLIN
+            )
+        ).generate()
+
+        val type = result.kotlinEnumTypes[0].members[0] as TypeSpec
+
+        assertThat(FileSpec.get("$basePackageName.enums", type).toString()).isEqualTo(
+            """
+                |package com.netflix.graphql.dgs.codegen.tests.generated.enums
+                |
+                |public enum class MyEnum {
+                |  A,
+                |  B,
+                |  C,
+                |  ;
+                |
+                |  public companion object
+                |}
+                |
+            """.trimMargin()
+        )
+
+        assertCompilesKotlin(result.kotlinEnumTypes)
+    }
 }