fix(tests): wait for KILL in TestExecution cleaner

Author: brian-mulier-p

tests/src/main/java/io/kestra/core/junit/extensions/KestraTestExtension.java

@@ -10,6 +10,7 @@
 import io.kestra.core.queues.TestQueueFactory;
 import io.kestra.core.repositories.ExecutionRepositoryInterface;
 import io.kestra.core.runners.TestRunner;
+import io.kestra.core.utils.Await;
 import io.kestra.core.utils.ListUtils;
 import io.kestra.core.utils.TestsUtils;
 import io.micronaut.inject.qualifiers.Qualifiers;
@@ -19,10 +20,12 @@
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.platform.commons.support.AnnotationSupport;
 
+import java.time.Duration;
 import java.util.Collections;
 import java.util.ConcurrentModificationException;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.TimeoutException;
 
 @Slf4j
 public class KestraTestExtension extends MicronautJunit5Extension {
@@ -81,9 +84,6 @@ public void afterTestExecution(ExtensionContext context) throws Exception {
 
         TestsUtils.queueConsumersCleanup();
 
-        KestraTest kestraTest = context.getTestClass()
-            .orElseThrow()
-            .getAnnotation(KestraTest.class);
         Optional<TestQueueFactory> testQueueFactory = Optional.of(applicationContext.containsBean(TestQueueFactory.class)).flatMap(contains -> contains ? Optional.of(applicationContext.getBean(TestQueueFactory.class)) : Optional.empty());
         List<Execution> testExecutions = testQueueFactory.map(TestQueueFactory::getTestExecutions).orElse(Collections.emptyList());
         if (!testExecutions.isEmpty()
@@ -100,12 +100,12 @@ public void afterTestExecution(ExtensionContext context) throws Exception {
 
 
     private void retryingExecutionKill(List<Execution> testExecutions, ExecutionRepositoryInterface executionRepository, QueueInterface<ExecutionKilled> killQueue, int retriesLeft) throws InterruptedException {
+        List<Execution> runningExecutions = ListUtils.distinctByKey(
+            testExecutions.stream().flatMap(launchedExecution -> executionRepository.findById(launchedExecution.getTenantId(), launchedExecution.getId()).stream()).toList(),
+            Execution::getId
+        ).stream().filter(inRepository -> !inRepository.getState().isTerminated()).toList();
         try {
-            ListUtils.distinctByKey(
-                    testExecutions.stream().flatMap(launchedExecution -> executionRepository.findById(launchedExecution.getTenantId(), launchedExecution.getId()).stream()).toList(),
-                    Execution::getId
-                ).stream().filter(inRepository -> inRepository.getState().isRunning() || inRepository.getState().isPaused() || inRepository.getState().isQueued())
-                .forEach(inRepository -> {
+            runningExecutions.forEach(inRepository -> {
                     log.warn("Execution {} is still running after test execution, killing it", inRepository.getId());
                     try {
                         killQueue.emit(ExecutionKilledExecution.builder()
@@ -128,5 +128,21 @@ private void retryingExecutionKill(List<Execution> testExecutions, ExecutionRepo
             Thread.sleep(100);
             retryingExecutionKill(testExecutions, executionRepository, killQueue, retriesLeft - 1);
         }
+
+        try {
+            Await.until(() -> runningExecutions.stream()
+                .map(execution -> executionRepository.findById(execution.getTenantId(), execution.getId()))
+                .allMatch(maybeExecution -> maybeExecution.map(inRepository -> {
+                        boolean terminated = inRepository.getState().isTerminated();
+                        if (!terminated) {
+                            log.warn("Execution {} has a pending KILL request but is still in state {} ", inRepository.getId(), inRepository.getState().getCurrent());
+                        }
+                        return terminated;
+                })
+                    .orElse(true))
+            , Duration.ofMillis(50), Duration.ofSeconds(10));
+        } catch (TimeoutException e) {
+            log.warn("Some executions remained in KILLING ", e);
+        }
     }
 }