Fix type resolvers for renamed polymorphic types and fix hooks signature (ExpediaGroup#282)

* devs can use @GraphQLName annotation to rename the underlying types - this change updates interface and union type resolvers to retrieve correct name when resolving the types
* update SchemaGeneratorHooks.didGenerateGraphQLType method to return the generated type instead of returning Unit, this makes it consistent with return type of other hook methods

Author: dariuszkuc

graphql-kotlin-schema-generator/src/main/kotlin/com/expedia/graphql/generator/TypeBuilder.kt

@@ -1,23 +1,26 @@
 package com.expedia.graphql.generator
 
+import com.expedia.graphql.SchemaGeneratorConfig
 import com.expedia.graphql.generator.extensions.getKClass
 import com.expedia.graphql.generator.extensions.isEnum
 import com.expedia.graphql.generator.extensions.isInterface
 import com.expedia.graphql.generator.extensions.isListType
 import com.expedia.graphql.generator.extensions.isUnion
 import com.expedia.graphql.generator.extensions.wrapInNonNull
 import com.expedia.graphql.generator.state.KGraphQLType
+import com.expedia.graphql.generator.state.SchemaGeneratorState
 import com.expedia.graphql.generator.state.TypesCacheKey
+import graphql.schema.GraphQLCodeRegistry
 import graphql.schema.GraphQLType
 import graphql.schema.GraphQLTypeReference
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
 
 internal open class TypeBuilder constructor(protected val generator: SchemaGenerator) {
-    protected val state = generator.state
-    protected val config = generator.config
-    protected val subTypeMapper = generator.subTypeMapper
-    protected val codeRegistry = generator.codeRegistry
+    protected val state: SchemaGeneratorState = generator.state
+    protected val config: SchemaGeneratorConfig = generator.config
+    protected val subTypeMapper: SubTypeMapper = generator.subTypeMapper
+    protected val codeRegistry: GraphQLCodeRegistry.Builder = generator.codeRegistry
 
     internal fun graphQLTypeOf(type: KType, inputType: Boolean = false, annotatedAsID: Boolean = false): GraphQLType {
         val hookGraphQLType = config.hooks.willGenerateGraphQLType(type)
@@ -26,10 +29,7 @@ internal open class TypeBuilder constructor(protected val generator: SchemaGener
             ?: objectFromReflection(type, inputType)
 
         val typeWithNullability = graphQLType.wrapInNonNull(type)
-
-        config.hooks.didGenerateGraphQLType(type, typeWithNullability)
-
-        return typeWithNullability
+        return config.hooks.didGenerateGraphQLType(type, typeWithNullability)
     }
 
     internal fun objectFromReflection(type: KType, inputType: Boolean): GraphQLType {

graphql-kotlin-schema-generator/src/main/kotlin/com/expedia/graphql/generator/types/InterfaceBuilder.kt

@@ -43,7 +43,7 @@ internal class InterfaceBuilder(generator: SchemaGenerator) : TypeBuilder(genera
                 }
             }
 
-            codeRegistry.typeResolver(interfaceType) { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObject<Any>().javaClass.simpleName) }
+            codeRegistry.typeResolver(interfaceType) { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObject<Any>().javaClass.kotlin.getSimpleName()) }
             config.hooks.onRewireGraphQLType(interfaceType).safeCast()
         }
     }

graphql-kotlin-schema-generator/src/main/kotlin/com/expedia/graphql/generator/types/UnionBuilder.kt

@@ -42,7 +42,7 @@ internal class UnionBuilder(generator: SchemaGenerator) : TypeBuilder(generator)
                 }
             }
             val unionType = builder.build()
-            codeRegistry.typeResolver(unionType) { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObject<Any>().javaClass.simpleName) }
+            codeRegistry.typeResolver(unionType) { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObject<Any>().javaClass.kotlin.getSimpleName()) }
             config.hooks.onRewireGraphQLType(unionType)
         }
     }

graphql-kotlin-schema-generator/src/main/kotlin/com/expedia/graphql/hooks/SchemaGeneratorHooks.kt

@@ -75,7 +75,7 @@ interface SchemaGeneratorHooks {
     /**
      * Called after wrapping the type based on nullity but before adding the generated type to the schema
      */
-    fun didGenerateGraphQLType(type: KType, generatedType: GraphQLType) = Unit
+    fun didGenerateGraphQLType(type: KType, generatedType: GraphQLType) = generatedType
 
     /**
      * Called after converting the function to a field definition but before adding to the schema to allow customization

graphql-kotlin-schema-generator/src/test/kotlin/com/expedia/graphql/generator/PolymorphicTests.kt

@@ -1,14 +1,18 @@
 package com.expedia.graphql.generator
 
 import com.expedia.graphql.TopLevelObject
+import com.expedia.graphql.annotations.GraphQLName
 import com.expedia.graphql.exceptions.InvalidInputFieldTypeException
 import com.expedia.graphql.testSchemaConfig
 import com.expedia.graphql.toSchema
+import graphql.TypeResolutionEnvironment
 import graphql.schema.GraphQLInterfaceType
 import graphql.schema.GraphQLObjectType
+import graphql.schema.GraphQLSchema
 import graphql.schema.GraphQLUnionType
 import org.junit.jupiter.api.Assertions.assertThrows
 import org.junit.jupiter.api.Test
+import kotlin.random.Random
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
@@ -93,6 +97,41 @@ internal class PolymorphicTests {
         assertEquals(1, classWithBaseAbstractType.interfaces.size)
         assertEquals(classWithBaseAbstractType.interfaces.first(), abstractInterface)
     }
+
+    @Test
+    fun `Interface types can be correctly resolved`() {
+        val schema = toSchema(queries = listOf(TopLevelObject(QueryWithRenamedAbstracts())), config = testSchemaConfig)
+
+        val cakeInterface = schema.getType("Cake") as? GraphQLInterfaceType
+        assertNotNull(cakeInterface)
+        val cakeResolver = schema.codeRegistry.getTypeResolver(cakeInterface)
+        val cheesecakeResolver = cakeResolver.getType(mockTypeResolutionEnvironment(Cheesecake(), schema))
+        assertNotNull(cheesecakeResolver)
+        assertEquals("Cheesecake", cheesecakeResolver.name)
+
+        val strawberryCakeResolver = cakeResolver.getType(mockTypeResolutionEnvironment(BerryCake(), schema))
+        assertNotNull(strawberryCakeResolver)
+        assertEquals("StrawberryCake", strawberryCakeResolver.name)
+    }
+
+    @Test
+    fun `Union types can be correctly resolved`() {
+        val schema = toSchema(queries = listOf(TopLevelObject(QueryWithRenamedAbstracts())), config = testSchemaConfig)
+
+        val dessertUnion = schema.getType("Dessert") as? GraphQLUnionType
+        assertNotNull(dessertUnion)
+        val dessertResolver = schema.codeRegistry.getTypeResolver(dessertUnion)
+        val iceCreamResolver = dessertResolver.getType(mockTypeResolutionEnvironment(IceCream(), schema))
+        assertNotNull(iceCreamResolver)
+        assertEquals("IceCream", iceCreamResolver.name)
+
+        val strawberryCakeResolver = dessertResolver.getType(mockTypeResolutionEnvironment(BerryCake(), schema))
+        assertNotNull(strawberryCakeResolver)
+        assertEquals("StrawberryCake", strawberryCakeResolver.name)
+    }
+
+    private fun mockTypeResolutionEnvironment(target: Any, schema: GraphQLSchema): TypeResolutionEnvironment =
+        TypeResolutionEnvironment(target, emptyMap(), null, null, schema, null)
 }
 
 class QueryWithInterface {
@@ -183,3 +222,35 @@ abstract class MyAbstract {
 }
 
 data class MyClass(override val id: Int, val name: String) : MyAbstract()
+
+class QueryWithRenamedAbstracts {
+
+    fun randomCake(): Cake = if (Random.nextBoolean()) {
+        BerryCake()
+    } else {
+        Cheesecake()
+    }
+
+    fun randomDessert(): Dessert = if (Random.nextBoolean()) {
+        IceCream()
+    } else {
+        BerryCake()
+    }
+}
+
+interface Cake {
+    fun recipe(): String
+}
+
+@GraphQLName("StrawberryCake")
+class BerryCake : Cake, Dessert {
+    override fun recipe(): String = "google it"
+}
+
+class Cheesecake : Cake {
+    override fun recipe(): String = "use bing"
+}
+
+interface Dessert
+
+class IceCream : Dessert

graphql-kotlin-schema-generator/src/test/kotlin/com/expedia/graphql/hooks/SchemaGeneratorHooksTest.kt

@@ -108,9 +108,10 @@ class SchemaGeneratorHooksTest {
         class MockSchemaGeneratorHooks : SchemaGeneratorHooks {
             var lastSeenType: KType? = null
             var lastSeenGeneratedType: GraphQLType? = null
-            override fun didGenerateGraphQLType(type: KType, generatedType: GraphQLType) {
+            override fun didGenerateGraphQLType(type: KType, generatedType: GraphQLType): GraphQLType {
                 lastSeenType = type
                 lastSeenGeneratedType = generatedType
+                return generatedType
             }
         }