Refactor AsyncQueryForwardingServletTest to reduce code duplication (#16092)

This commit is contained in:
gzhao9 2024-03-10 08:02:43 -04:00 committed by GitHub
parent ba3d4daf45
commit 2d628cce84
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -278,28 +278,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
public void testHandleExceptionWithFilterDisabled() throws Exception public void testHandleExceptionWithFilterDisabled() throws Exception
{ {
String errorMessage = "test exception message"; 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.assertTrue(captor.getValue() instanceof QueryException);
Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode()); Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
Assert.assertEquals(errorMessage, captor.getValue().getMessage()); Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@ -310,41 +294,25 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
public void testHandleExceptionWithFilterEnabled() throws Exception public void testHandleExceptionWithFilterEnabled() throws Exception
{ {
String errorMessage = "test exception message"; 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); @Override
Mockito.when(response.getOutputStream()).thenReturn(outputStream); public boolean isShowDetailedJettyErrors()
final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( {
new MapQueryToolChestWarehouse(ImmutableMap.of()), return true;
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;
}
@Override @Override
public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() public ErrorResponseTransformStrategy getErrorResponseTransformStrategy()
{ {
return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); 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.assertTrue(captor.getValue() instanceof QueryException);
Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode()); Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
Assert.assertNull(captor.getValue().getMessage()); Assert.assertNull(captor.getValue().getMessage());
@ -356,41 +324,25 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
public void testHandleExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception public void testHandleExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception
{ {
String errorMessage = "test exception message"; 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); @Override
Mockito.when(response.getOutputStream()).thenReturn(outputStream); public boolean isShowDetailedJettyErrors()
final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( {
new MapQueryToolChestWarehouse(ImmutableMap.of()), return true;
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;
}
@Override @Override
public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() public ErrorResponseTransformStrategy getErrorResponseTransformStrategy()
{ {
return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*")); 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.assertTrue(captor.getValue() instanceof QueryException);
Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode()); Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
Assert.assertEquals(errorMessage, captor.getValue().getMessage()); Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@ -402,30 +354,12 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
public void testHandleQueryParseExceptionWithFilterDisabled() throws Exception public void testHandleQueryParseExceptionWithFilterDisabled() throws Exception
{ {
String errorMessage = "test exception message"; 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.assertTrue(captor.getValue() instanceof QueryException);
Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode()); Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
Assert.assertEquals(errorMessage, captor.getValue().getMessage()); Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@ -436,43 +370,24 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
public void testHandleQueryParseExceptionWithFilterEnabled() throws Exception public void testHandleQueryParseExceptionWithFilterEnabled() throws Exception
{ {
String errorMessage = "test exception message"; String errorMessage = "test exception message";
ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class); ServerConfig serverConfig = new ServerConfig() {
HttpServletRequest request = Mockito.mock(HttpServletRequest.class); @Override
HttpServletResponse response = Mockito.mock(HttpServletResponse.class); public boolean isShowDetailedJettyErrors()
ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class); {
Mockito.when(response.getOutputStream()).thenReturn(outputStream); return true;
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;
}
@Override @Override
public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() public ErrorResponseTransformStrategy getErrorResponseTransformStrategy()
{ {
return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); 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.assertTrue(captor.getValue() instanceof QueryException);
Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode()); Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
Assert.assertNull(captor.getValue().getMessage()); Assert.assertNull(captor.getValue().getMessage());
@ -484,43 +399,25 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
public void testHandleQueryParseExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception public void testHandleQueryParseExceptionWithFilterEnabledButMessageMatchAllowedRegex() throws Exception
{ {
String errorMessage = "test exception message"; 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); @Override
ServletOutputStream outputStream = Mockito.mock(ServletOutputStream.class); public boolean isShowDetailedJettyErrors()
Mockito.when(response.getOutputStream()).thenReturn(outputStream); {
final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( return true;
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;
}
@Override @Override
public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() public ErrorResponseTransformStrategy getErrorResponseTransformStrategy()
{ {
return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of("test .*")); 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.assertTrue(captor.getValue() instanceof QueryException);
Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode()); Assert.assertEquals("Unknown exception", ((QueryException) captor.getValue()).getErrorCode());
Assert.assertEquals(errorMessage, captor.getValue().getMessage()); Assert.assertEquals(errorMessage, captor.getValue().getMessage());
@ -1063,4 +960,36 @@ public class AsyncQueryForwardingServletTest extends BaseJettyTest
return port; 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;
}
} }