### Question:
RetryOnRenew implements HttpRetryHandler { @Override public boolean shouldRetryRequest(HttpCommand command, HttpResponse response) { boolean retry = false; switch (response.getStatusCode()) { case 401: Multimap<String, String> headers = command.getCurrentRequest().getHeaders(); if (headers != null && headers.containsKey(AuthHeaders.AUTH_USER) && headers.containsKey(AuthHeaders.AUTH_KEY) && !headers.containsKey(AuthHeaders.AUTH_TOKEN)) { retry = false; } else { Integer count = retryCountMap.getIfPresent(command); if (count == null) { logger.debug("invalidating authentication token - first time for %s", command); retryCountMap.put(command, 1); authenticationResponseCache.invalidateAll(); retry = true; } else { if (count + 1 >= NUM_RETRIES) { logger.debug("too many 401s - giving up after: %s for %s", count, command); retry = false; } else { logger.debug("invalidating authentication token - retry %s for %s", count, command); retryCountMap.put(command, count + 1); authenticationResponseCache.invalidateAll(); Uninterruptibles.sleepUninterruptibly(5, TimeUnit.SECONDS); retry = true; } } } break; } return retry; } @Inject protected  RetryOnRenew(LoadingCache<Credentials, AuthenticationResponse> authenticationResponseCache); @Override boolean shouldRetryRequest(HttpCommand command, HttpResponse response);  }### Answer:
@Test public void test401ShouldRetry() { HttpCommand command = createMock(HttpCommand.class); HttpRequest request = createMock(HttpRequest.class); HttpResponse response = createMock(HttpResponse.class); @SuppressWarnings("unchecked") LoadingCache<Credentials, AuthenticationResponse> cache = createMock(LoadingCache.class); expect(command.getCurrentRequest()).andReturn(request); cache.invalidateAll(); expectLastCall(); expect(response.getPayload()).andReturn(Payloads.newStringPayload("")) .anyTimes(); expect(response.getStatusCode()).andReturn(401).atLeastOnce(); replay(command); replay(response); replay(cache); RetryOnRenew retry = new RetryOnRenew(cache); assertTrue(retry.shouldRetryRequest(command, response)); verify(command); verify(response); verify(cache); }
@Test public void test401ShouldRetry4Times() { HttpCommand command = createMock(HttpCommand.class); HttpRequest request = createMock(HttpRequest.class); HttpResponse response = createMock(HttpResponse.class); @SuppressWarnings("unchecked") LoadingCache<Credentials, AuthenticationResponse> cache = createMock(LoadingCache.class); expect(command.getCurrentRequest()).andReturn(request).anyTimes(); expect(request.getHeaders()).andStubReturn(null); cache.invalidateAll(); expectLastCall().anyTimes(); expect(response.getPayload()).andReturn(Payloads.newStringPayload("")) .anyTimes(); expect(response.getStatusCode()).andReturn(401).anyTimes(); replay(command, request, response, cache); RetryOnRenew retry = new RetryOnRenew(cache); for (int n = 0; n < RetryOnRenew.NUM_RETRIES - 1; n++) { assertTrue(retry.shouldRetryRequest(command, response), "Expected retry to succeed"); } assertFalse(retry.shouldRetryRequest(command, response), "Expected retry to fail on attempt 5"); verify(command, response, cache); }