Refactor AsyncQueryForwardingServletTest to reduce code duplication (#16092)

services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java

@@ -278,28 +278,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
   public void testHandleExceptionWithFilterDisabled() throws Exception
   {
     String errorMessage = "test exception message";
-    ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+    ServerConfig serverConfig = new ServerConfig();
-    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
+    ArgumentCaptor<Exception> captor = captureExceptionHandledByServlet(
-    ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class);
+        serverConfig,
-    Mockito.when(response.getOutputStream()).thenReturn(outputStream);
+        (servlet, request, response, mapper)
-    final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet(
+            -> servlet.handleException(response, mapper, new IllegalStateException(errorMessage))
-        new MapQueryToolChestWarehouse(ImmutableMap.of()),
-        mockMapper,
-        TestHelper.makeSmileMapper(),
-        null,
-        null,
-        null,
-        new NoopServiceEmitter(),
-        new NoopRequestLogger(),
-        new DefaultGenericQueryMetricsFactory(),
-        new AuthenticatorMapper(ImmutableMap.of()),
-        new Properties(),
-        new ServerConfig()
     );
-    Exception testException = new IllegalStateException(errorMessage);
-    servlet.handleException(response, mockMapper, testException);
-    ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(Exception.class);
-    Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture());
     Assert.assertTrue(captor.getValue() instanceof QueryException);
     Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
     Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@@ -310,23 +294,7 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
   public void testHandleExceptionWithFilterEnabled() throws Exception
   {
     String errorMessage = "test exception message";
-    ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+    ServerConfig serverConfig = new ServerConfig() 
-    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
-    ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class);
-    Mockito.when(response.getOutputStream()).thenReturn(outputStream);
-    final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet(
-        new MapQueryToolChestWarehouse(ImmutableMap.of()),
-        mockMapper,
-        TestHelper.makeSmileMapper(),
-        null,
-        null,
-        null,
-        new NoopServiceEmitter(),
-        new NoopRequestLogger(),
-        new DefaultGenericQueryMetricsFactory(),
-        new AuthenticatorMapper(ImmutableMap.of()),
-        new Properties(),
-        new ServerConfig()
     {
       @Override
       public boolean isShowDetailedJettyErrors()
@@ -339,12 +307,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
       {
         return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of());
       }
-        }
+    };
+    ArgumentCaptor<Exception> captor = captureExceptionHandledByServlet(
+        serverConfig,
+        (servlet, request, response, mapper)
+            -> servlet.handleException(response, mapper, new IllegalStateException(errorMessage))
     );
-    Exception testException = new IllegalStateException(errorMessage);
-    servlet.handleException(response, mockMapper, testException);
-    ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(Exception.class);
-    Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture());
     Assert.assertTrue(captor.getValue() instanceof QueryException);
     Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
     Assert.assertNull(captor.getValue().getMessage());
@@ -356,23 +324,7 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
   public void testHandleExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception
   {
     String errorMessage = "test exception message";
-    ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+    ServerConfig serverConfig = new ServerConfig()
-    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
-    ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class);
-    Mockito.when(response.getOutputStream()).thenReturn(outputStream);
-    final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet(
-        new MapQueryToolChestWarehouse(ImmutableMap.of()),
-        mockMapper,
-        TestHelper.makeSmileMapper(),
-        null,
-        null,
-        null,
-        new NoopServiceEmitter(),
-        new NoopRequestLogger(),
-        new DefaultGenericQueryMetricsFactory(),
-        new AuthenticatorMapper(ImmutableMap.of()),
-        new Properties(),
-        new ServerConfig()
     {
       @Override
       public boolean isShowDetailedJettyErrors()
@@ -385,12 +337,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
       {
         return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*"));
       }
-        }
+    };
+    ArgumentCaptor<Exception> captor = captureExceptionHandledByServlet(
+        serverConfig,
+        (servlet, request, response, mapper)
+            -> servlet.handleException(response, mapper, new IllegalStateException(errorMessage))
     );
-    Exception testException = new IllegalStateException(errorMessage);
-    servlet.handleException(response, mockMapper, testException);
-    ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(Exception.class);
-    Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture());
     Assert.assertTrue(captor.getValue() instanceof QueryException);
     Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
     Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@@ -402,30 +354,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
   public void testHandleQueryParseExceptionWithFilterDisabled() throws Exception
   {
     String errorMessage = "test exception message";
-    ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+    ServerConfig serverConfig = new ServerConfig();
-    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
+    ArgumentCaptor<Exception> captor = captureExceptionHandledByServlet(
-    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
+        serverConfig,
-    ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class);
+        (servlet, request, response, mapper)
-    Mockito.when(response.getOutputStream()).thenReturn(outputStream);
+            -> servlet.handleQueryParseException(request, response, mapper, new IOException(errorMessage), false)
-    final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet(
-        new MapQueryToolChestWarehouse(ImmutableMap.of()),
-        mockMapper,
-        TestHelper.makeSmileMapper(),
-        null,
-        null,
-        null,
-        new NoopServiceEmitter(),
-        new NoopRequestLogger(),
-        new DefaultGenericQueryMetricsFactory(),
-        new AuthenticatorMapper(ImmutableMap.of()),
-        new Properties(),
-        new ServerConfig()
     );
-    Mockito.when(request.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).thenReturn(new AuthenticationResult("userA", "basic", "basic", null));
-    IOException testException = new IOException(errorMessage);
-    servlet.handleQueryParseException(request, response, mockMapper, testException, false);
-    ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(Exception.class);
-    Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture());
     Assert.assertTrue(captor.getValue() instanceof QueryException);
     Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
     Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@@ -436,25 +370,7 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
   public void testHandleQueryParseExceptionWithFilterEnabled() throws Exception
   {
     String errorMessage = "test exception message";
-    ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+    ServerConfig serverConfig = new ServerConfig() {
-    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
-    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
-    ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class);
-    Mockito.when(response.getOutputStream()).thenReturn(outputStream);
-    final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet(
-        new MapQueryToolChestWarehouse(ImmutableMap.of()),
-        mockMapper,
-        TestHelper.makeSmileMapper(),
-        null,
-        null,
-        null,
-        new NoopServiceEmitter(),
-        new NoopRequestLogger(),
-        new DefaultGenericQueryMetricsFactory(),
-        new AuthenticatorMapper(ImmutableMap.of()),
-        new Properties(),
-        new ServerConfig()
-        {
       @Override
       public boolean isShowDetailedJettyErrors()
       {
@@ -466,13 +382,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
       {
         return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of());
       }
-        }
+    };
+    ArgumentCaptor<Exception> captor = captureExceptionHandledByServlet(
+        serverConfig,
+        (servlet, request, response, mapper)
+            -> servlet.handleQueryParseException(request, response, mapper, new IOException(errorMessage), false)
     );
-    Mockito.when(request.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).thenReturn(new AuthenticationResult("userA", "basic", "basic", null));
-    IOException testException = new IOException(errorMessage);
-    servlet.handleQueryParseException(request, response, mockMapper, testException, false);
-    ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(Exception.class);
-    Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture());
     Assert.assertTrue(captor.getValue() instanceof QueryException);
     Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
     Assert.assertNull(captor.getValue().getMessage());
@@ -484,24 +399,7 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
   public void testHandleQueryParseExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception
   {
     String errorMessage = "test exception message";
-    ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+    ServerConfig serverConfig = new ServerConfig()
-    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
-    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
-    ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class);
-    Mockito.when(response.getOutputStream()).thenReturn(outputStream);
-    final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet(
-        new MapQueryToolChestWarehouse(ImmutableMap.of()),
-        mockMapper,
-        TestHelper.makeSmileMapper(),
-        null,
-        null,
-        null,
-        new NoopServiceEmitter(),
-        new NoopRequestLogger(),
-        new DefaultGenericQueryMetricsFactory(),
-        new AuthenticatorMapper(ImmutableMap.of()),
-        new Properties(),
-        new ServerConfig()
     {
       @Override
       public boolean isShowDetailedJettyErrors()
@@ -514,13 +412,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
       {
         return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*"));
       }
-        }
+    };
+    ArgumentCaptor<Exception> captor = captureExceptionHandledByServlet(
+        serverConfig,
+        (servlet, request, response, mapper)
+            -> servlet.handleQueryParseException(request, response, mapper, new IOException(errorMessage), false)
     );
-    Mockito.when(request.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).thenReturn(new AuthenticationResult("userA", "basic", "basic", null));
-    IOException testException = new IOException(errorMessage);
-    servlet.handleQueryParseException(request, response, mockMapper, testException, false);
-    ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(Exception.class);
-    Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture());
     Assert.assertTrue(captor.getValue() instanceof QueryException);
     Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
     Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@@ -1063,4 +960,36 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
       return port;
     }
   }
+  interface ServletTestAction
+  {
+    void execute(AsyncQueryForwardingServlet servlet, HttpServletRequest request, HttpServletResponse response, ObjectMapper mapper) throws Exception;
+  }
+
+  private ArgumentCaptor<Exception> captureExceptionHandledByServlet(ServerConfig serverConfig, ServletTestAction action) throws Exception
+  {
+    ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+    HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
+    ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class);
+    Mockito.when(response.getOutputStream()).thenReturn(outputStream);
+    final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet(
+            new MapQueryToolChestWarehouse(ImmutableMap.of()),
+            mockMapper,
+            TestHelper.makeSmileMapper(),
+            null,
+            null,
+            null,
+            new NoopServiceEmitter(),
+            new NoopRequestLogger(),
+            new DefaultGenericQueryMetricsFactory(),
+            new AuthenticatorMapper(ImmutableMap.of()),
+            new Properties(),
+            serverConfig
+    );
+    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
+    Mockito.when(request.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).thenReturn(new AuthenticationResult("userA", "basic", "basic", null));
+    action.execute(servlet, request, response, mockMapper);
+    ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(Exception.class);
+    Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture());
+    return captor;
+  }
 }