Resolve union and interface values by the raw class of a parameterized type

src/main/kotlin/graphql/kickstart/tools/DictionaryTypeResolver.kt

@@ -3,18 +3,29 @@ package graphql.kickstart.tools
 import graphql.TypeResolutionEnvironment
 import graphql.kickstart.tools.util.BiMap
 import graphql.kickstart.tools.util.JavaType
+import graphql.kickstart.tools.util.unwrap
 import graphql.language.TypeDefinition
 import graphql.schema.GraphQLInterfaceType
 import graphql.schema.GraphQLObjectType
 import graphql.schema.GraphQLUnionType
 import graphql.schema.TypeResolver
+import java.lang.reflect.ParameterizedType
 
 /**
  * @author Andrew Potter
  */
 internal abstract class DictionaryTypeResolver(
         private val dictionary: BiMap<JavaType, TypeDefinition<*>>
 ) : TypeResolver {
+    // Raw class of each parameterized entry, e.g. Box for Box<String>, unless several entries share it
+    private val rawClassDictionary by lazy {
+        dictionary.entries
+            .filter { it.key is ParameterizedType }
+            .groupBy({ it.key.unwrap() }, { it.value })
+            .mapNotNull { (clazz, definitions) -> definitions.singleOrNull()?.let { clazz to it } }
+            .toMap()
+    }
+
     private fun <T> getTypeDefinition(clazz: Class<T>): TypeDefinition<*>? {
         return dictionary[clazz]
             ?: (if (clazz.superclass == null) null else getTypeDefinition(clazz.superclass))
@@ -24,7 +35,11 @@ internal abstract class DictionaryTypeResolver(
     override fun getType(env: TypeResolutionEnvironment): GraphQLObjectType? {
         val clazz = env.getObject<Any>()!!.javaClass
         val name = getTypeDefinition(clazz)?.name ?: clazz.simpleName
-        return env.schema.getObjectType(name) ?: throw TypeResolverError(getError(name))
+        return env.schema.getType(name) as? GraphQLObjectType
+            ?: rawClassDictionary[clazz]?.let { env.schema.getType(it.name) as? GraphQLObjectType }
+            // throws graphql-java's AssertException when the name is a type other than an object type
+            ?: env.schema.getObjectType(name)
+            ?: throw TypeResolverError(getError(name))
     }
 
     abstract fun getError(name: String): String

src/test/kotlin/graphql/kickstart/tools/GenericObjectTypesTest.kt

@@ -1,6 +1,7 @@
 package graphql.kickstart.tools
 
 import graphql.GraphQL
+import org.junit.Assert.assertThrows
 import org.junit.Test
 
 class GenericObjectTypesTest {
@@ -12,8 +13,11 @@ class GenericObjectTypesTest {
                 """
                 type Query {
                     test: TestClass!
+                    anyMatch: Match!
                 }
 
+                union Match = StringListEntityMatch | StringEntityMatch
+
                 type TestClass {
                     match: StringListEntityMatch!
                     single: StringEntityMatch!
@@ -32,7 +36,7 @@ class GenericObjectTypesTest {
             .resolvers(EntityMatchQueryResolver())
             .build()
             .makeExecutableSchema()
-            .assertEveryFieldResolves()
+            .assertEveryFieldResolves(except = setOf("Query.anyMatch"))
         val gql = GraphQL.newGraphQL(schema).build()
 
         val data = assertNoGraphQlErrors(gql) {
@@ -50,6 +54,9 @@ class GenericObjectTypesTest {
             "match" to mapOf("text" to "x y", "value" to listOf("x", "y")),
             "single" to mapOf("text" to "z", "value" to "z")
         ))
+
+        val error = assertThrows(TypeResolverError::class.java) { gql.execute("query { anyMatch { __typename } }") }
+        assertEquals(error.message, "Expected object type with name 'EntityMatch' to exist for union 'Match', but it doesn't!")
     }
 
     @Test
@@ -60,14 +67,22 @@ class GenericObjectTypesTest {
                 type Query {
                     characterHit: CharacterHit!
                     characterHitPage: CharacterHitPage!
+                    lastPage: Result!
+                    topHit: Scored!
+                }
+
+                union Result = CharacterHitPage | Human
+
+                interface Scored {
+                    score: Float!
                 }
 
                 type CharacterHitPage {
                     content: [CharacterHit!]!
                     totalElements: Int!
                 }
 
-                type CharacterHit {
+                type CharacterHit implements Scored {
                     score: Float!
                     source: Character!
                 }
@@ -110,6 +125,13 @@ class GenericObjectTypesTest {
                         }
                     }
                 }
+                lastPage {
+                    ... on CharacterHitPage { totalElements }
+                }
+                topHit {
+                    score
+                    ... on CharacterHit { source { name } }
+                }
             }
             """
         }
@@ -122,13 +144,56 @@ class GenericObjectTypesTest {
                 mapOf("source" to mapOf("name" to "R2-D2", "primaryFunction" to "Astromech"))
             )
         ))
+        assertEquals(data["lastPage"], mapOf("totalElements" to 1))
+        assertEquals(data["topHit"], mapOf("score" to 0.5, "source" to mapOf("name" to "Luke")))
+    }
+
+    @Test
+    fun `a raw class keeps its own type when a parameterized entry shares it`() {
+        val schema = SchemaParser.newParser()
+            .schemaString(
+                """
+                type Query {
+                    single: StringEntityMatch!
+                    anyMatch: Match!
+                }
+
+                union Match = StringEntityMatch | AnyEntityMatch
+
+                type StringEntityMatch {
+                    text: String!
+                    value: String
+                }
+
+                type AnyEntityMatch {
+                    text: String!
+                }
+                """)
+            .resolvers(RawEntityMatchQueryResolver())
+            .dictionary("AnyEntityMatch", EntityMatch::class.java)
+            .build()
+            .makeExecutableSchema()
+            .assertEveryFieldResolves()
+        val gql = GraphQL.newGraphQL(schema).build()
+
+        val data = assertNoGraphQlErrors(gql) { "query { anyMatch { __typename } }" }
+
+        assertEquals(data["anyMatch"], mapOf("__typename" to "AnyEntityMatch"))
     }
 
     class EntityMatchQueryResolver : GraphQLQueryResolver {
         fun test(): TestClass = TestClass(
             match = EntityMatch("x y", listOf("x", "y")),
             single = EntityMatch("z", "z")
         )
+
+        fun anyMatch(): Any = EntityMatch("z", "z")
+    }
+
+    class RawEntityMatchQueryResolver : GraphQLQueryResolver {
+        fun single(): EntityMatch<String> = EntityMatch("z", "z")
+
+        fun anyMatch(): Any = EntityMatch("z", "z")
     }
 
     data class TestClass(val match: EntityMatch<List<String>>, val single: EntityMatch<String>)
@@ -141,11 +206,19 @@ class GenericObjectTypesTest {
             listOf(Hit(Human("Luke", "Tatooine"), 1.0), Hit(Droid("R2-D2", "Astromech"), 1.0)),
             2
         )
+
+        fun lastPage(): Any = Page(listOf(Hit(Droid("R2-D2", "Astromech"), 1.0)), 1)
+
+        fun topHit(): Scored = Hit(Human("Luke", "Tatooine"), 0.5)
     }
 
     class Page<T>(val content: List<T>, val totalElements: Int)
 
-    class Hit<T>(val source: T, val score: Double)
+    interface Scored {
+        val score: Double
+    }
+
+    class Hit<T>(val source: T, override val score: Double) : Scored
 
     interface Character {
         val name: String