fix(tests): remove ThreadLocal + add distinct on TestExecution cleaner

Author: brian-mulier-p

core/src/main/java/io/kestra/core/utils/ListUtils.java

@@ -2,6 +2,10 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+import java.util.function.Predicate;
 
 public class ListUtils {
     public static <T> List<T> emptyOnNull(List<T> list) {
@@ -71,4 +75,13 @@ public static List<String> convertToListString(Object object){
             throw new IllegalArgumentException("%s in not an instance of List of String".formatted(object));
         }
     }
+
+    private static <T> Predicate<T> distinctByKeyPredicate(Function<? super T,Object> keyExtractor) {
+        Map<Object,Boolean> seen = new ConcurrentHashMap<>();
+        return t -> seen.putIfAbsent(keyExtractor.apply(t), Boolean.TRUE) == null;
+    }
+
+    public static <T> List<T> distinctByKey(List<T> withDuplicates, Function<? super T, Object> keyExtractor) {
+        return withDuplicates.stream().filter(distinctByKeyPredicate(keyExtractor)).toList();
+    }
 }

core/src/test/java/io/kestra/core/utils/ListUtilsTest.java

@@ -1,7 +1,11 @@
 package io.kestra.core.utils;
 
+import io.kestra.core.models.executions.Execution;
+import io.kestra.core.models.flows.State;
 import org.junit.jupiter.api.Test;
 
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Collections;
 import java.util.List;
 
@@ -52,4 +56,19 @@ void convertToListString(){
         assertThrows(IllegalArgumentException.class, () -> ListUtils.convertToListString("not a list"));
         assertThrows(IllegalArgumentException.class, () -> ListUtils.convertToListString(List.of(1, 2, 3)));
     }
-}
+
+    @Test
+    void distinctByKey() {
+        List<Execution> executions = List.of(
+            Execution.builder().id("1").flowId("flow1").state(State.of(State.Type.CREATED, Collections.emptyList())).build(),
+            Execution.builder().id("1").flowId("flow1").state(State.of(State.Type.RUNNING, List.of(new State.History(State.Type.CREATED, Instant.now().minus(Duration.ofSeconds(2)))))).build(),
+            Execution.builder().id("2").flowId("flow2").build()
+        );
+
+        List<Execution> distinctExecutions = ListUtils.distinctByKey(executions, Execution::getId);
+
+        assertThat(distinctExecutions.size()).isEqualTo(2);
+        assertThat(distinctExecutions.stream().map(Execution::getId)).containsExactlyInAnyOrder("1", "2");
+        assertThat(distinctExecutions.stream().filter(e -> e.getId().equals("1")).findFirst().get().getState().getCurrent()).isEqualTo(State.Type.CREATED);
+    }
+}

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

@@ -10,24 +10,22 @@
 import io.kestra.core.queues.TestQueueFactory;
 import io.kestra.core.repositories.ExecutionRepositoryInterface;
 import io.kestra.core.runners.TestRunner;
+import io.kestra.core.utils.ListUtils;
 import io.kestra.core.utils.TestsUtils;
-import io.micronaut.core.annotation.NonNull;
-import io.micronaut.inject.BeanDefinition;
 import io.micronaut.inject.qualifiers.Qualifiers;
 import io.micronaut.test.annotation.MicronautTestValue;
 import io.micronaut.test.extensions.junit5.MicronautJunit5Extension;
 import lombok.extern.slf4j.Slf4j;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.platform.commons.support.AnnotationSupport;
 
-import java.util.ArrayList;
+import java.util.Collections;
 import java.util.ConcurrentModificationException;
 import java.util.List;
+import java.util.Optional;
 
 @Slf4j
 public class KestraTestExtension extends MicronautJunit5Extension {
-    public static final ThreadLocal<List<Execution>> testExecutions = ThreadLocal.withInitial(ArrayList::new);
-
     private static final ExtensionContext.Namespace NAMESPACE = ExtensionContext.Namespace.create(KestraTestExtension.class);
 
     @Override
@@ -77,14 +75,6 @@ public void beforeAll(ExtensionContext extensionContext) throws Exception {
         }
     }
 
-    @Override
-    public void beforeTestExecution(ExtensionContext context) {
-        if (applicationContext.containsBean(TestQueueFactory.class)) {
-            TestQueueFactory testQueueFactory = applicationContext.getBean(TestQueueFactory.class);
-            testQueueFactory.setTestExecutionsList(testExecutions.get());
-        }
-    }
-
     @Override
     public void afterTestExecution(ExtensionContext context) throws Exception {
         super.afterTestExecution(context);
@@ -94,24 +84,27 @@ public void afterTestExecution(ExtensionContext context) throws Exception {
         KestraTest kestraTest = context.getTestClass()
             .orElseThrow()
             .getAnnotation(KestraTest.class);
-        if (!testExecutions.get().isEmpty() &&
-            kestraTest.startRunner()
+        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()
             && applicationContext.containsBean(ExecutionRepositoryInterface.class)
             && applicationContext.containsBean(QueueInterface.class, Qualifiers.byName(QueueFactoryInterface.KILL_NAMED))) {
             ExecutionRepositoryInterface executionRepository = applicationContext.getBean(ExecutionRepositoryInterface.class);
             QueueInterface<ExecutionKilled> killQueue = applicationContext.getBean(QueueInterface.class, Qualifiers.byName(QueueFactoryInterface.KILL_NAMED));
 
-            retryingExecutionKill(executionRepository, killQueue, 10);
+            retryingExecutionKill(testExecutions, executionRepository, killQueue, 10);
 
-            testExecutions.get().clear();
+            testExecutions.clear();
         }
     }
 
-    private void retryingExecutionKill(ExecutionRepositoryInterface executionRepository, QueueInterface<ExecutionKilled> killQueue, int retriesLeft) throws InterruptedException {
+
+    private void retryingExecutionKill(List<Execution> testExecutions, ExecutionRepositoryInterface executionRepository, QueueInterface<ExecutionKilled> killQueue, int retriesLeft) throws InterruptedException {
         try {
-            testExecutions.get().stream()
-                .flatMap(launchedExecution -> executionRepository.findById(launchedExecution.getTenantId(), launchedExecution.getId()).stream())
-                .filter(inRepository -> inRepository.getState().isRunning() || inRepository.getState().isPaused() || inRepository.getState().isQueued())
+            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 -> {
                     log.warn("Execution {} is still running after test execution, killing it", inRepository.getId());
                     try {
@@ -133,7 +126,7 @@ private void retryingExecutionKill(ExecutionRepositoryInterface executionReposit
                 return;
             }
             Thread.sleep(100);
-            retryingExecutionKill(executionRepository, killQueue, retriesLeft - 1);
+            retryingExecutionKill(testExecutions, executionRepository, killQueue, retriesLeft - 1);
         }
     }
 }

tests/src/main/java/io/kestra/core/queues/TestQueueFactory.java

@@ -5,8 +5,11 @@
 import io.micronaut.core.annotation.Introspected;
 import jakarta.inject.Named;
 import jakarta.inject.Singleton;
+import lombok.Getter;
+import lombok.Setter;
 
 import java.lang.reflect.Proxy;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
@@ -15,16 +18,13 @@
 @Requires(bean = QueueFactoryInterface.class)
 public class TestQueueFactory {
     private QueueInterface<Execution> delegate;
-    private List<Execution> testExecutions;
+    @Getter
+    private List<Execution> testExecutions = new ArrayList<>();
 
     public TestQueueFactory(QueueFactoryInterface queueFactoryInterface) {
         this.delegate = queueFactoryInterface.execution();
     }
 
-    public void setTestExecutionsList(List<Execution> testExecutions) {
-        this.testExecutions = testExecutions;
-    }
-
     @SuppressWarnings("unchecked")
     @Singleton
     @Replaces(named = QueueFactoryInterface.EXECUTION_NAMED)