Fix package name for type mapped union sub types (#731)

Netflix · Aug 16, 2024 · 69f01a1 · 69f01a1
1 parent 76679a6
commit 69f01a1
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 3 deletions.
diff --git a/graphql-dgs-codegen-core/src/main/kotlin/com/netflix/graphql/dgs/codegen/CodeGen.kt b/graphql-dgs-codegen-core/src/main/kotlin/com/netflix/graphql/dgs/codegen/CodeGen.kt
@@ -355,7 +355,7 @@ class CodeGen(private val config: CodeGenConfig) {
                 .excludeSchemaTypeExtension()
                 .map {
                     val extensions = findUnionExtensions(it.name, definitions)
-                    KotlinUnionTypeGenerator(config).generate(it, extensions)
+                    KotlinUnionTypeGenerator(config, document).generate(it, extensions)
                 }
                 .fold(CodeGenResult.EMPTY) { result, next -> result.merge(next) }
 

diff --git a/...main/kotlin/com/netflix/graphql/dgs/codegen/generators/kotlin/KotlinUnionTypeGenerator.kt b/...main/kotlin/com/netflix/graphql/dgs/codegen/generators/kotlin/KotlinUnionTypeGenerator.kt
@@ -24,13 +24,15 @@ import com.netflix.graphql.dgs.codegen.shouldSkip
 import com.squareup.kotlinpoet.ClassName
 import com.squareup.kotlinpoet.FileSpec
 import com.squareup.kotlinpoet.TypeSpec
+import graphql.language.Document
 import graphql.language.TypeName
 import graphql.language.UnionTypeDefinition
 import graphql.language.UnionTypeExtensionDefinition
 
-class KotlinUnionTypeGenerator(private val config: CodeGenConfig) {
+class KotlinUnionTypeGenerator(private val config: CodeGenConfig, document: Document) {
 
     private val packageName = config.packageNameTypes
+    private val typeUtils = KotlinTypeUtils(packageName, config, document)
 
     fun generate(definition: UnionTypeDefinition, extensions: List<UnionTypeExtensionDefinition>): CodeGenResult {
         if (definition.shouldSkip(config)) {
@@ -42,7 +44,8 @@ class KotlinUnionTypeGenerator(private val config: CodeGenConfig) {
 
         val memberTypes = definition.memberTypes.plus(extensions.flatMap { it.memberTypes }).asSequence()
             .filterIsInstance<TypeName>()
-            .map { member -> ClassName(packageName, member.name) }
+            .map { member -> typeUtils.findKtInterfaceName(member.name, packageName) }
+            .filterIsInstance<ClassName>()
             .toList()
 
         if (memberTypes.isNotEmpty()) {

diff --git a/...hql-dgs-codegen-core/src/test/kotlin/com/netflix/graphql/dgs/codegen/KotlinCodeGenTest.kt b/...hql-dgs-codegen-core/src/test/kotlin/com/netflix/graphql/dgs/codegen/KotlinCodeGenTest.kt
@@ -1114,6 +1114,56 @@ class KotlinCodeGenTest {
         )
     }
 
+    @Test
+    fun `Use mapped type name when the type implements non-mapped union type`() {
+        val schema = """
+            union Pet = Cat | Dog
+            type Cat implements Pet {
+              name: ID!
+            }
+            type Dog implements Pet {
+              name: ID!
+            }
+        """.trimIndent()
+
+        val codeGenResult = CodeGen(
+            CodeGenConfig(
+                schemas = setOf(schema),
+                packageName = basePackageName,
+                language = Language.KOTLIN,
+                typeMapping = mapOf(
+                    "Cat" to "mypackage.Cat"
+                )
+            )
+        ).generate()
+        val interfaces = codeGenResult.kotlinInterfaces
+
+        assertThat(interfaces.size).isEqualTo(1)
+        assertThat(interfaces[0].toString()).isEqualTo(
+            """
+                |package com.netflix.graphql.dgs.codegen.tests.generated.types
+                |
+                |import com.fasterxml.jackson.`annotation`.JsonSubTypes
+                |import com.fasterxml.jackson.`annotation`.JsonTypeInfo
+                |import mypackage.Cat
+                |
+                |@JsonTypeInfo(
+                |  use = JsonTypeInfo.Id.NAME,
+                |  include = JsonTypeInfo.As.PROPERTY,
+                |  property = "__typename",
+                |)
+                |@JsonSubTypes(value = [
+                |  JsonSubTypes.Type(value = Cat::class, name = "Cat"),
+                |  JsonSubTypes.Type(value = Dog::class, name = "Dog")
+                |])
+                |public interface Pet {
+                |  public companion object
+                |}
+                |
+            """.trimMargin()
+        )
+    }
+
     class MappedTypesTestCases : ArgumentsProvider {
         override fun provideArguments(context: ExtensionContext): Stream<out Arguments> = of(
             arguments("java.time.LocalDateTime", "java.time.LocalDateTime"),