fix(DataTypeGenerator): support typeMapping for union parent type (#552)

Netflix · May 11, 2023 · 51969ec · 51969ec
1 parent ffef604
commit 51969ec
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 1 deletion.
diff --git a/...core/src/main/kotlin/com/netflix/graphql/dgs/codegen/generators/java/DataTypeGenerator.kt b/...core/src/main/kotlin/com/netflix/graphql/dgs/codegen/generators/java/DataTypeGenerator.kt
@@ -394,7 +394,10 @@ abstract class BaseDataTypeGenerator(
     }
 
     private fun addInterface(type: String, javaType: TypeSpec.Builder) {
-        javaType.addSuperinterface(ClassName.get(packageName, type))
+        val interfaceTypeMappedName: String? = config.typeMapping[type]
+        val interfaceName: ClassName = if (interfaceTypeMappedName == null) ClassName.get(packageName, type) else ClassName.bestGuess(interfaceTypeMappedName)
+
+        javaType.addSuperinterface(interfaceName)
     }
 
     private fun addField(fieldDefinition: Field, javaType: TypeSpec.Builder) {

diff --git a/graphql-dgs-codegen-core/src/test/kotlin/com/netflix/graphql/dgs/codegen/CodeGenTest.kt b/graphql-dgs-codegen-core/src/test/kotlin/com/netflix/graphql/dgs/codegen/CodeGenTest.kt
@@ -4151,4 +4151,34 @@ It takes a title and such.
         // Check that the third field of the Person type is an Integer
         assertThat(dataTypes[0].typeSpec.fieldSpecs[2].type.toString()).isEqualTo("java.lang.Integer")
     }
+
+    @Test
+    fun `Supports typeMapping for union types`() {
+        val schema = """
+            type A {
+                name: String
+            }
+        
+            type B {
+                count: Int
+            }
+        
+            union C = A | B
+        """.trimIndent()
+
+        val (dataTypes) = CodeGen(
+            CodeGenConfig(
+                schemas = setOf(schema),
+                packageName = basePackageName,
+                typeMapping = mapOf(
+                    "C" to "java.lang.String"
+                )
+            )
+        ).generate()
+
+        assertThat(dataTypes.size).isEqualTo(2)
+
+        assertThat(dataTypes[0].typeSpec.superinterfaces[0].toString()).isEqualTo("java.lang.String")
+        assertThat(dataTypes[1].typeSpec.superinterfaces[0].toString()).isEqualTo("java.lang.String")
+    }
 }