diff --git a/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java b/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java
index 472fa6c96f2d..bd39d8f6add7 100644
--- a/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java
+++ b/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java
@@ -30,6 +30,7 @@
 import org.apache.druid.java.util.common.concurrent.Execs;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryCapacityExceededException;
+import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryException;
 import org.apache.druid.query.QueryInterruptedException;
 import org.apache.druid.query.QueryTimeoutException;
@@ -306,10 +307,15 @@ public void testTimeoutAfterCallingFuture()
     private Query<?> mockQuery(String queryId, long timeoutAt)
     {
       Query<?> query = Mockito.mock(Query.class);
+      QueryContext context = Mockito.mock(QueryContext.class);
       Mockito.when(query.getId()).thenReturn(queryId);
-      Mockito.when(query.getQueryContext().getAsLong(ArgumentMatchers.eq(DirectDruidClient.QUERY_FAIL_TIME),
-                                                     ArgumentMatchers.eq(-1L)))
-             .thenReturn(timeoutAt);
+      Mockito.when(query.getQueryContext()).thenReturn(context);
+      Mockito.when(
+          context.getAsLong(
+              ArgumentMatchers.eq(DirectDruidClient.QUERY_FAIL_TIME),
+              ArgumentMatchers.eq(-1L)
+          )
+      ).thenReturn(timeoutAt);
       return query;
     }
   }