Fail the field when the custom context class is missing from the GraphQLContext

Author: oryan-block

src/main/kotlin/graphql/kickstart/tools/resolver/MethodFieldResolver.kt

@@ -102,9 +102,17 @@ internal class MethodFieldResolver(
         val firstTrailingParameter = field.inputValueDefinitions.size + search.firstArgumentIndex
         method.parameterTypes.slice(firstTrailingParameter until numberOfParameters).forEach { parameterType ->
             when (parameterType) {
-                options.contextClass -> args.add { environment -> environment.graphQlContext[options.contextClass] }
-
                 GraphQLContext::class.java -> args.add { environment -> environment.graphQlContext }
+
+                options.contextClass -> args.add { environment ->
+                    environment.graphQlContext.get<Any>(parameterType) ?: throw ResolverError(
+                        "Resolver method '${method.declaringClass.name}.${method.name}' has a parameter of the context class " +
+                            "'${parameterType.name}', but the GraphQLContext has no value for it. Put it in the GraphQLContext under " +
+                            "the class key when building the ExecutionInput, for example " +
+                            "graphQLContext(mapOf(${parameterType.simpleName}::class.java to context))."
+                    )
+                }
+
                 else -> args.add { environment -> environment }
             }
         }

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

@@ -19,6 +19,8 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.channels.ReceiveChannel
+import org.junit.Assert.assertThrows
+import org.junit.Assert.assertTrue
 import org.junit.Test
 import org.reactivestreams.Publisher
 import org.reactivestreams.tck.TestEnvironment
@@ -276,15 +278,44 @@ class MethodFieldResolverDataFetcherTest {
     }
 
     @Test
-    fun `data fetcher does not read custom context from the legacy context object`() {
+    fun `data fetcher fails if custom context is missing from GraphQLContext`() {
         val options = SchemaParserOptions.newOptions().contextClass(ContextClass::class).build()
         val resolver = createFetcher("active", options = options, resolver = object : GraphQLResolver<DataClass> {
             fun isActive(dataClass: DataClass, ctx: ContextClass?): Boolean {
                 return ctx != null
             }
         })
 
-        assertEquals(resolver.get(createEnvironment(DataClass(), legacyContext = ContextClass())), false)
+        val error = assertThrows(ResolverError::class.java) { resolver.get(createEnvironment(DataClass())) }
+
+        assertTrue(error.message!!.contains("'${ContextClass::class.java.name}'"))
+        assertTrue(error.message!!.contains(".isActive'"))
+        assertTrue(error.message!!.contains("GraphQLContext"))
+    }
+
+    @Test
+    fun `data fetcher fails if custom context is only passed as the legacy context object`() {
+        val options = SchemaParserOptions.newOptions().contextClass(ContextClass::class).build()
+        val resolver = createFetcher("active", options = options, resolver = object : GraphQLResolver<DataClass> {
+            fun isActive(dataClass: DataClass, ctx: ContextClass): Boolean {
+                return true
+            }
+        })
+
+        assertThrows(ResolverError::class.java) { resolver.get(createEnvironment(DataClass(), legacyContext = ContextClass())) }
+    }
+
+    @Test
+    fun `data fetcher passes GraphQLContext if it is the custom context class`() {
+        val context = GraphQLContext.newContext().build()
+        val options = SchemaParserOptions.newOptions().contextClass(GraphQLContext::class).build()
+        val resolver = createFetcher("active", options = options, resolver = object : GraphQLResolver<DataClass> {
+            fun isActive(dataClass: DataClass, ctx: GraphQLContext): Boolean {
+                return ctx == context
+            }
+        })
+
+        assertEquals(resolver.get(createEnvironment(DataClass(), context = context)), true)
     }
 
     @Test