@Test
public void testLoadBalancingWithFailure() throws Exception {

  Configuration conf = new Configuration();
  KMSClientProvider p1 = mock(KMSClientProvider.class);
  when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
      .thenReturn(
          new KMSClientProvider.KMSKeyVersion("p1", "v1", new byte[0]));
  when(p1.getKMSUrl()).thenReturn("p1");
  // This should not be retried
  KMSClientProvider p2 = mock(KMSClientProvider.class);
  when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
      .thenThrow(new NoSuchAlgorithmException("p2"));
  when(p2.getKMSUrl()).thenReturn("p2");
  KMSClientProvider p3 = mock(KMSClientProvider.class);
  when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
      .thenReturn(
          new KMSClientProvider.KMSKeyVersion("p3", "v3", new byte[0]));
  when(p3.getKMSUrl()).thenReturn("p3");
  // This should be retried
  KMSClientProvider p4 = mock(KMSClientProvider.class);
  when(p4.createKey(Mockito.anyString(), Mockito.any(Options.class)))
      .thenThrow(new IOException("p4"));
  when(p4.getKMSUrl()).thenReturn("p4");
  KeyProvider kp = new LoadBalancingKMSClientProvider(
      new KMSClientProvider[] { p1, p2, p3, p4 }, 0, conf);
  assertEquals("p1", kp.createKey("test4", new Options(conf)).getName());
  // Exceptions other than IOExceptions will not be retried
  try {
    kp.createKey("test1", new Options(conf)).getName();
    fail("Should fail since its not an IOException");
  } catch (Exception e) {
    assertTrue(e instanceof NoSuchAlgorithmException);
  }
  assertEquals("p3", kp.createKey("test2", new Options(conf)).getName());
  // IOException will trigger retry in next provider
  assertEquals("p1", kp.createKey("test3", new Options(conf)).getName());
}