From 2d628cce848450def64200b7448cc987279d1d32 Mon Sep 17 00:00:00 2001 From: gzhao9 <74684732+gzhao9@users.noreply.github.com> Date: Sun, 10 Mar 2024 08:02:43 -0400 Subject: [PATCH] Refactor AsyncQueryForwardingServletTest to reduce code duplication (#16092) --- .../AsyncQueryForwardingServletTest.java | 289 +++++++----------- 1 file changed, 109 insertions(+), 180 deletions(-) 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 034a0563453..29831f3c430 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 class AsyncQueryForwardingServletTest extends BaseJettyTest 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 class AsyncQueryForwardingServletTest extends BaseJettyTest 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 class AsyncQueryForwardingServletTest extends BaseJettyTest 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 class AsyncQueryForwardingServletTest extends BaseJettyTest 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 class AsyncQueryForwardingServletTest extends BaseJettyTest 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 class AsyncQueryForwardingServletTest extends BaseJettyTest 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 class AsyncQueryForwardingServletTest extends BaseJettyTest 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; + } }