- final String msg = "Expected instance of NoShardLeaderException, actual: \n"
- + Throwables.getStackTraceAsString(e.getCause());
- assertTrue(msg, Throwables.getRootCause(e) instanceof NoShardLeaderException);
- assertEquals(DistributedDataStore.class, testParameter);
- } catch (TimeoutException e) {
- // ClientBackedDataStore doesn't set cause to ExecutionException, future just time outs
- assertEquals(ClientBackedDataStore.class, testParameter);
+ final String msg = "Unexpected exception: " + Throwables.getStackTraceAsString(e.getCause());
+ if (DistributedDataStore.class.equals(testParameter)) {
+ assertTrue(msg, Throwables.getRootCause(e) instanceof NoShardLeaderException);
+ } else {
+ assertTrue(msg, Throwables.getRootCause(e) instanceof RequestTimeoutException);
+ }