@Test
public void handle_WTEWithCasThrown_ShouldThrowProperExecutionException()
    throws ExecutionException {
  // Arrange
  put = preparePutWithClusteringKey();
  put.withCondition(new PutIfNotExists());
  spy = prepareSpiedInsertStatementHandler();

  WriteTimeoutException toThrow = mock(WriteTimeoutException.class);
  when(toThrow.getWriteType()).thenReturn(WriteType.CAS);
  doThrow(toThrow).when(spy).handleInternal(put);

  // Act Assert
  assertThatThrownBy(
          () -> {
            spy.handle(put);
          })
      .isInstanceOf(RetriableExecutionException.class)
      .hasCause(toThrow);
}
 

@Test
public void
    handle_PutWithConditionGivenAndWTEWithSimpleThrown_ShouldThrowProperExecutionException()
        throws ExecutionException {
  // Arrange
  put = preparePutWithClusteringKey();
  put.withCondition(new PutIfNotExists());
  spy = prepareSpiedInsertStatementHandler();

  WriteTimeoutException toThrow = mock(WriteTimeoutException.class);
  when(toThrow.getWriteType()).thenReturn(WriteType.SIMPLE);
  doThrow(toThrow).when(spy).handleInternal(put);

  // Act Assert
  assertThatThrownBy(
          () -> {
            spy.handle(put);
          })
      .isInstanceOf(ReadRepairableExecutionException.class)
      .hasCause(toThrow);
}
 

@Test
public void
    handle_PutWithoutConditionGivenAndWTEWithSimpleThrown_ShouldThrowProperExecutionException()
        throws ExecutionException {
  // Arrange
  put = preparePutWithClusteringKey();
  spy = prepareSpiedInsertStatementHandler();

  WriteTimeoutException toThrow = mock(WriteTimeoutException.class);
  when(toThrow.getWriteType()).thenReturn(WriteType.SIMPLE);
  doThrow(toThrow).when(spy).handleInternal(put);

  // Act Assert
  assertThatThrownBy(
          () -> {
            spy.handle(put);
          })
      .isInstanceOf(RetriableExecutionException.class)
      .hasCause(toThrow);
}
 

@Test
public void handle_WTEThrownInLoggingInBatchExecution_ShouldThrowRetriableExecutionException() {
  // Arrange
  configureBehavior();
  mutations = prepareConditionalPuts();
  WriteTimeoutException e = mock(WriteTimeoutException.class);
  when(e.getWriteType()).thenReturn(WriteType.BATCH_LOG);
  when(session.execute(any(Statement.class))).thenThrow(e);

  // Act Assert
  assertThatThrownBy(
          () -> {
            batch.handle(mutations);
          })
      .isInstanceOf(RetriableExecutionException.class)
      .hasCause(e);
}
 

@Test
public void handle_WTEThrownInMutationInBatchExecution_ShouldExecuteProperly() {
  // Arrange
  configureBehavior();
  mutations = prepareConditionalPuts();
  WriteTimeoutException e = mock(WriteTimeoutException.class);
  when(e.getWriteType()).thenReturn(WriteType.BATCH);
  when(session.execute(any(Statement.class))).thenThrow(e);

  // Act Assert
  assertThatCode(
          () -> {
            batch.handle(mutations);
          })
      .doesNotThrowAnyException();
}
 

@Test
public void handle_WTEThrownInCasInBatchExecution_ShouldThrowRetriableExecutionException() {
  // Arrange
  configureBehavior();
  mutations = prepareConditionalPuts();
  WriteTimeoutException e = mock(WriteTimeoutException.class);
  when(e.getWriteType()).thenReturn(WriteType.CAS);
  when(session.execute(any(Statement.class))).thenThrow(e);

  // Act Assert
  assertThatThrownBy(
          () -> {
            batch.handle(mutations);
          })
      .isInstanceOf(RetriableExecutionException.class)
      .hasCause(e);
}
 

@Test
public void
    handle_WTEThrownInSimpleWriteInBatchExecution_ShouldThrowRetriableExecutionException() {
  // Arrange
  configureBehavior();
  mutations = prepareConditionalPuts();
  WriteTimeoutException e = mock(WriteTimeoutException.class);
  when(e.getWriteType()).thenReturn(WriteType.SIMPLE);
  when(session.execute(any(Statement.class))).thenThrow(e);

  // Act Assert
  assertThatThrownBy(
          () -> {
            batch.handle(mutations);
          })
      .isInstanceOf(RetriableExecutionException.class)
      .hasCause(e);
}
 

@Test
public void testShouldReturnWriteTimeout() throws Exception {
  server.prime(when(query).then(writeTimeout(ConsistencyLevel.QUORUM, 3, 1, WriteType.BATCH)));

  thrown.expect(WriteTimeoutException.class);
  thrown.expect(
      match(
          (WriteTimeoutException wte) ->
              wte.getRequiredAcknowledgements() == 1
                  && wte.getReceivedAcknowledgements() == 3
                  && wte.getConsistencyLevel() == com.datastax.driver.core.ConsistencyLevel.QUORUM
                  && wte.getWriteType() == com.datastax.driver.core.WriteType.BATCH));

  query();
}
 

protected void checkCassandraException(Exception e) {
    _mexceptions.incr();
    if (e instanceof AlreadyExistsException ||
            e instanceof AuthenticationException ||
            e instanceof DriverException ||
            e instanceof DriverInternalError ||
            e instanceof InvalidConfigurationInQueryException ||
            e instanceof InvalidQueryException ||
            e instanceof InvalidTypeException ||
            e instanceof QueryExecutionException ||
            e instanceof QueryValidationException ||
            e instanceof ReadTimeoutException ||
            e instanceof SyntaxError ||
            e instanceof TraceRetrievalException ||
            e instanceof TruncateException ||
            e instanceof UnauthorizedException ||
            e instanceof UnavailableException ||
            e instanceof ReadTimeoutException ||
            e instanceof WriteTimeoutException ||
            e instanceof ReadFailureException ||
            e instanceof WriteFailureException ||
            e instanceof FunctionExecutionException) {
        throw new ReportedFailedException(e);
    } else {
        throw new RuntimeException(e);
    }
}