diff --git a/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java b/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java index 034a05634534..29831f3c4302 100644 --- a/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java +++ b/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java @@ -278,28 +278,12 @@ public void testJDBCSqlProxy() throws Exception public void testHandleExceptionWithFilterDisabled() throws Exception { String errorMessage = "test exception message"; - 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(), - new ServerConfig() + ServerConfig serverConfig = new ServerConfig(); + ArgumentCaptor 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 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,41 +294,25 @@ public void testHandleExceptionWithFilterDisabled() throws Exception public void testHandleExceptionWithFilterEnabled() throws Exception { String errorMessage = "test exception message"; - 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(), - new ServerConfig() - { - @Override - public boolean isShowDetailedJettyErrors() - { - return true; - } + ServerConfig serverConfig = new ServerConfig() + { + @Override + public boolean isShowDetailedJettyErrors() + { + return true; + } - @Override - public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() - { - return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); - } - } + @Override + public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() + { + return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); + } + }; + ArgumentCaptor 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 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,41 +324,25 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() public void testHandleExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception { String errorMessage = "test exception message"; - 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(), - new ServerConfig() - { - @Override - public boolean isShowDetailedJettyErrors() - { - return true; - } + ServerConfig serverConfig = new ServerConfig() + { + @Override + public boolean isShowDetailedJettyErrors() + { + return true; + } - @Override - public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() - { - return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*")); - } - } + @Override + public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() + { + return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*")); + } + }; + ArgumentCaptor 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 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 ErrorResponseTransformStrategy getErrorResponseTransformStrategy() public void testHandleQueryParseExceptionWithFilterDisabled() throws Exception { String errorMessage = "test exception message"; - ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class); - 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() + ServerConfig serverConfig = new ServerConfig(); + ArgumentCaptor 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 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,43 +370,24 @@ public void testHandleQueryParseExceptionWithFilterDisabled() throws Exception public void testHandleQueryParseExceptionWithFilterEnabled() throws Exception { String errorMessage = "test exception message"; - ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class); - 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() - { - return true; - } + ServerConfig serverConfig = new ServerConfig() { + @Override + public boolean isShowDetailedJettyErrors() + { + return true; + } - @Override - public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() - { - return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); - } - } + @Override + public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() + { + return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); + } + }; + ArgumentCaptor 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 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,43 +399,25 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() public void testHandleQueryParseExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception { String errorMessage = "test exception message"; - ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class); - 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() - { - return true; - } + ServerConfig serverConfig = new ServerConfig() + { + @Override + public boolean isShowDetailedJettyErrors() + { + return true; + } - @Override - public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() - { - return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*")); - } - } + @Override + public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() + { + return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*")); + } + }; + ArgumentCaptor 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 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 int getPort() return port; } } + interface ServletTestAction + { + void execute(AsyncQueryForwardingServlet servlet, HttpServletRequest request, HttpServletResponse response, ObjectMapper mapper) throws Exception; + } + + private ArgumentCaptor 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 captor = ArgumentCaptor.forClass(Exception.class); + Mockito.verify(mockMapper).writeValue(ArgumentMatchers.eq(outputStream), captor.capture()); + return captor; + } }