refactor test code

.

Signed-off-by: moonyougnCHAE <[email protected]>

Author: moonyoungCHAE

spring-kafka/src/test/java/org/springframework/kafka/listener/TransactionalContainerTests.java

@@ -76,6 +76,8 @@
 import org.springframework.kafka.transaction.KafkaTransactionManager;
 import org.springframework.messaging.MessageHeaders;
 import org.springframework.transaction.TransactionDefinition;
+import org.springframework.transaction.TransactionExecution;
+import org.springframework.transaction.TransactionExecutionListener;
 import org.springframework.transaction.annotation.Transactional;
 import org.springframework.transaction.support.DefaultTransactionDefinition;
 import org.springframework.util.backoff.FixedBackOff;
@@ -1167,7 +1169,6 @@ void testSendOffsetOnlyOnActiveTransaction() throws InterruptedException {
 		DefaultKafkaConsumerFactory<Integer, String> cf = new DefaultKafkaConsumerFactory<>(consumerProperties);
 		ContainerProperties containerProps = new ContainerProperties(topic11);
 		containerProps.setPollTimeout(10_000);
-		final var successLatch = new AtomicReference<>(new CountDownLatch(2));
 		containerProps.setMessageListener(new MessageListener<Integer, String>() {
 			@Transactional("testSendOffsetOnlyOnActiveTransaction")
 			@Override
@@ -1177,9 +1178,19 @@ public void onMessage(ConsumerRecord<Integer, String> data) {
 
 		// init container
 		KafkaTransactionManager<Object, Object> tm = new KafkaTransactionManager<>(pf);
+		AtomicInteger txCount = new AtomicInteger(0);
+		tm.addListener(new TransactionExecutionListener() {
+			@Override
+			public void afterCommit(TransactionExecution transaction, @Nullable Throwable commitFailure) {
+				txCount.incrementAndGet();
+				TransactionExecutionListener.super.afterCommit(transaction, commitFailure);
+			}
+		});
 		containerProps.setKafkaAwareTransactionManager(tm);
+
 		KafkaMessageListenerContainer<Integer, String> container = new KafkaMessageListenerContainer<>(cf, containerProps);
 		container.setBeanName("testSendOffsetOnlyOnActiveTransaction");
+		final var interceptorLatch = new AtomicReference<>(new CountDownLatch(1));
 		container.setRecordInterceptor(new RecordInterceptor<Integer, String>() {
 			boolean isFirst = true;
 
@@ -1198,17 +1209,25 @@ public void onMessage(ConsumerRecord<Integer, String> data) {
 			public void afterRecord(
 					ConsumerRecord<Integer, String> record,
 					Consumer<Integer, String> consumer) {
-				successLatch.get().countDown();
+				interceptorLatch.get().countDown();
 			}
 		});
 		container.start();
 
 		template.executeInTransaction(t -> {
 			template.send(new ProducerRecord<>(topic11, 0, 0, "bar1"));
+			return null;
+		});
+		assertThat(interceptorLatch.get().await(30, TimeUnit.SECONDS)).isTrue();
+		assertThat(txCount.get()).isEqualTo(1);
+
+		interceptorLatch.set(new CountDownLatch(1));
+		template.executeInTransaction(t -> {
 			template.send(new ProducerRecord<>(topic11, 0, 0, "bar2"));
 			return null;
 		});
-		assertThat(successLatch.get().await(30, TimeUnit.SECONDS)).isTrue();
+		assertThat(interceptorLatch.get().await(30, TimeUnit.SECONDS)).isTrue();
+		assertThat(txCount.get()).isEqualTo(1);
 
 		container.stop();
 		pf.destroy();