Source:
<code>public void testDropCommitOnBatchExpiry() throws InterruptedException { final long pid = 13131L; final short epoch = 1; doInitTransactions(pid, epoch); transactionManager.beginTransaction(); transactionManager.failIfNotReadyForSend(); transactionManager.maybeAddPartitionToTransaction(tp0); Future<RecordMetadata> responseFuture = appendToAccumulator(tp0); assertFalse(responseFuture.isDone()); prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid); assertFalse(transactionManager.transactionContainsPartition(tp0)); assertFalse(transactionManager.isSendToPartitionAllowed(tp0)); runUntil(() -> transactionManager.transactionContainsPartition(tp0)); assertTrue(transactionManager.isSendToPartitionAllowed(tp0)); assertFalse(responseFuture.isDone()); TransactionalRequestResult commitResult = transactionManager.beginCommit(); time.sleep(10000); Node clusterNode = metadata.fetch().nodes().get(0); client.disconnect(clusterNode.idString()); runUntil(responseFuture::isDone); try { responseFuture.get(); fail("Expected to get a TimeoutException since the queued ProducerBatch should have been expired"); } catch (ExecutionException e) { assertTrue(e.getCause() instanceof TimeoutException); } runUntil(commitResult::isCompleted); assertFalse(commitResult.isSuccessful()); assertTrue(transactionManager.hasAbortableError()); assertTrue(transactionManager.hasOngoingTransaction()); assertFalse(transactionManager.isCompleting()); assertTrue(transactionManager.transactionContainsPartition(tp0)); TransactionalRequestResult abortResult = transactionManager.beginAbort(); prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch); <START> sender.runOnce(); END> prepareInitPidResponse(Errors.NONE, false, pid, (short) (epoch + 1)); sender.runOnce(); runUntil(abortResult::isCompleted); assertTrue(abortResult.isSuccessful()); assertFalse(transactionManager.hasOngoingTransaction()); assertFalse(transactionManager.transactionContainsPartition(tp0)); }</code><technical_language>call runUntil, get rid of individual calls runOnce. For example: java prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch); prepareInitPidResponse(Errors.NONE, false, pid, (short) (epoch + 1)); runUntil(abortResult::isCompleted); Similarly for other test cases</technical_language>

Target:
public void testDropCommitOnBatchExpiry() throws InterruptedException { final long pid = 13131L; final short epoch = 1; doInitTransactions(pid, epoch); transactionManager.beginTransaction(); transactionManager.failIfNotReadyForSend(); transactionManager.maybeAddPartitionToTransaction(tp0); Future<RecordMetadata> responseFuture = appendToAccumulator(tp0); assertFalse(responseFuture.isDone()); prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid); assertFalse(transactionManager.transactionContainsPartition(tp0)); assertFalse(transactionManager.isSendToPartitionAllowed(tp0)); runUntil(() -> transactionManager.transactionContainsPartition(tp0)); assertTrue(transactionManager.isSendToPartitionAllowed(tp0)); assertFalse(responseFuture.isDone()); TransactionalRequestResult commitResult = transactionManager.beginCommit(); time.sleep(10000); Node clusterNode = metadata.fetch().nodes().get(0); client.disconnect(clusterNode.idString()); runUntil(responseFuture::isDone); try { responseFuture.get(); fail("Expected to get a TimeoutException since the queued ProducerBatch should have been expired"); } catch (ExecutionException e) { assertTrue(e.getCause() instanceof TimeoutException); } runUntil(commitResult::isCompleted); assertFalse(commitResult.isSuccessful()); assertTrue(transactionManager.hasAbortableError()); assertTrue(transactionManager.hasOngoingTransaction()); assertFalse(transactionManager.isCompleting()); assertTrue(transactionManager.transactionContainsPartition(tp0)); TransactionalRequestResult abortResult = transactionManager.beginAbort(); prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch); prepareInitPidResponse(Errors.NONE, false, pid, (short) (epoch + 1)); runUntil(abortResult::isCompleted); assertTrue(abortResult.isSuccessful()); assertFalse(transactionManager.hasOngoingTransaction()); assertFalse(transactionManager.transactionContainsPartition(tp0)); }