From f4ec50bf7a9e7e423dad4b1c2fb3c5af846ef075 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Tue, 13 Sep 2022 20:49:57 -0700 Subject: [PATCH] fix JsonParserIteratorTest (#13083) --- .../apache/druid/client/JsonParserIteratorTest.java | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) 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 472fa6c96f2..bd39d8f6add 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.StringUtils; 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 class JsonParserIteratorTest 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; } }