From caccfdb7c22e81edb602d7fff8127f39c686ed79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Mathieu?= <loikeseke@gmail.com>
Date: Wed, 5 Mar 2025 17:41:33 +0100
Subject: [PATCH] WIP

---
 .../java/io/kestra/core/models/Label.java     |   1 +
 .../core/models/executions/Execution.java     |  28 ++++-
 .../kestra/core/runners/ExecutorService.java  |  31 +++++-
 .../java/io/kestra/core/test/Fixture.java     |  23 ++++
 .../main/java/io/kestra/core/test/Input.java  |  15 +++
 .../main/java/io/kestra/core/test/Test.java   |  24 +++++
 test1.yaml                                    |  21 ++++
 test2.yaml                                    |  13 +++
 .../controllers/api/ExecutionController.java  |   6 +-
 .../controllers/api/TestController.java       | 100 ++++++++++++++++++
 10 files changed, 254 insertions(+), 8 deletions(-)
 create mode 100644 core/src/main/java/io/kestra/core/test/Fixture.java
 create mode 100644 core/src/main/java/io/kestra/core/test/Input.java
 create mode 100644 core/src/main/java/io/kestra/core/test/Test.java
 create mode 100644 test1.yaml
 create mode 100644 test2.yaml
 create mode 100644 webserver/src/main/java/io/kestra/webserver/controllers/api/TestController.java

diff --git a/core/src/main/java/io/kestra/core/models/Label.java b/core/src/main/java/io/kestra/core/models/Label.java
index 4ca80a49201..99c6cd3c610 100644
--- a/core/src/main/java/io/kestra/core/models/Label.java
+++ b/core/src/main/java/io/kestra/core/models/Label.java
@@ -18,6 +18,7 @@ public record Label(@NotNull String key, @NotNull String value) {
     public static final String RESTARTED = SYSTEM_PREFIX + "restarted";
     public static final String REPLAY = SYSTEM_PREFIX + "replay";
     public static final String REPLAYED = SYSTEM_PREFIX + "replayed";
+    public static final String TEST = SYSTEM_PREFIX + "test";
 
     /**
      * Static helper method for converting a list of labels to a nested map.
diff --git a/core/src/main/java/io/kestra/core/models/executions/Execution.java b/core/src/main/java/io/kestra/core/models/executions/Execution.java
index 4696b28f3fe..af8d2c8ef58 100644
--- a/core/src/main/java/io/kestra/core/models/executions/Execution.java
+++ b/core/src/main/java/io/kestra/core/models/executions/Execution.java
@@ -21,7 +21,9 @@
 import io.kestra.core.serializers.ListOrMapOfLabelDeserializer;
 import io.kestra.core.serializers.ListOrMapOfLabelSerializer;
 import io.kestra.core.services.LabelService;
+import io.kestra.core.test.Fixture;
 import io.kestra.core.utils.IdUtils;
+import io.kestra.core.utils.ListUtils;
 import io.kestra.core.utils.MapUtils;
 import io.micronaut.core.annotation.Nullable;
 import io.swagger.v3.oas.annotations.Hidden;
@@ -111,6 +113,10 @@ public class Execution implements DeletedInterface, TenantInterface {
     @Setter
     String traceParent;
 
+    @With
+    @Nullable
+    List<Fixture> fixtures;
+
     /**
      * Factory method for constructing a new {@link Execution} object for the given {@link Flow}.
      *
@@ -205,7 +211,8 @@ public Execution withState(State.Type state) {
             this.deleted,
             this.metadata,
             this.scheduleDate,
-            this.traceParent
+            this.traceParent,
+            this.fixtures
         );
     }
 
@@ -229,7 +236,8 @@ public Execution withLabels(List<Label> labels) {
             this.deleted,
             this.metadata,
             this.scheduleDate,
-            this.traceParent
+            this.traceParent,
+            this.fixtures
         );
     }
 
@@ -266,7 +274,8 @@ public Execution withTaskRun(TaskRun taskRun) throws InternalException {
             this.deleted,
             this.metadata,
             this.scheduleDate,
-            this.traceParent
+            this.traceParent,
+            this.fixtures
         );
     }
 
@@ -290,7 +299,8 @@ public Execution childExecution(String childExecutionId, List<TaskRun> taskRunLi
             this.deleted,
             this.metadata,
             this.scheduleDate,
-            this.traceParent
+            this.traceParent,
+            this.fixtures
         );
     }
 
@@ -723,6 +733,16 @@ public FailedExecutionWithLog failedExecutionFromExecutor(Exception e) {
             );
     }
 
+    public Optional<Fixture> getFixtureForTaskRun(TaskRun taskRun) {
+        if (this.fixtures == null) {
+            return Optional.empty();
+        }
+
+        return this.fixtures.stream()
+            .filter(fixture -> Objects.equals(fixture.getId(), taskRun.getTaskId()) && Objects.equals(fixture.getValue(), taskRun.getValue()))
+            .findFirst();
+    }
+
     /**
      * Create a new attempt for failed worker execution
      *
diff --git a/core/src/main/java/io/kestra/core/runners/ExecutorService.java b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
index 7c576176557..04ab29d9459 100644
--- a/core/src/main/java/io/kestra/core/runners/ExecutorService.java
+++ b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
@@ -15,7 +15,9 @@
 import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.queues.QueueInterface;
 import io.kestra.core.services.*;
+import io.kestra.core.test.Fixture;
 import io.kestra.core.trace.propagation.RunContextTextMapSetter;
+import io.kestra.core.utils.IdUtils;
 import io.kestra.core.utils.ListUtils;
 import io.kestra.core.utils.TruthUtils;
 import io.kestra.plugin.core.flow.Pause;
@@ -759,7 +761,7 @@ private Executor handleWorkerTask(final Executor executor) throws InternalExcept
         Map<Boolean, List<WorkerTask>> workerTasks = executor.getExecution()
             .getTaskRunList()
             .stream()
-            .filter(taskRun -> taskRun.getState().getCurrent().isCreated())
+            .filter(taskRun -> taskRun.getState().getCurrent().isCreated() && executor.getExecution().getFixtureForTaskRun(taskRun).isEmpty())
             .map(throwFunction(taskRun -> {
                     Task task = executor.getFlow().findTaskByTaskId(taskRun.getTaskId());
                     RunContext runContext = runContextFactory.of(executor.getFlow(), task, executor.getExecution(), taskRun);
@@ -812,7 +814,32 @@ private Executor handleWorkerTask(final Executor executor) throws InternalExcept
             )
             .collect(Collectors.groupingBy(workerTask -> workerTask.getTaskRun().getState().isFailed() || workerTask.getTaskRun().getState().getCurrent() == State.Type.CANCELLED));
 
-        if (workerTasks.isEmpty()) {
+        // mock WorkerTaskResult for mocked execution
+        // submit TaskRun when receiving created, must be done after the state execution store
+        // TODO check if it cannot be done directly inside the Execution
+        boolean hasMockedWorkerTask = false;
+        record FixtureAndTaskRun(Fixture fixture, TaskRun taskRun) {} // FIXME temporal
+        if (executor.getExecution().getFixtures() != null) {
+            List<WorkerTaskResult> workerTaskResults = executor.getExecution()
+                .getTaskRunList()
+                .stream()
+                .filter(taskRun -> taskRun.getState().getCurrent().isCreated())
+                .flatMap(taskRun -> executor.getExecution().getFixtureForTaskRun(taskRun).stream().map(fixture -> new FixtureAndTaskRun(fixture, taskRun)))
+                .map(fixtureAndTaskRun -> WorkerTaskResult.builder()
+                    .taskRun(fixtureAndTaskRun.taskRun()
+                        .withState(Optional.ofNullable(fixtureAndTaskRun.fixture().getState()).orElse(State.Type.SUCCESS))
+                        .withOutputs(fixtureAndTaskRun.fixture().getOutputs())
+                    )
+                    .build()
+                )
+                .toList();
+
+            hasMockedWorkerTask = !workerTaskResults.isEmpty();
+            this.addWorkerTaskResults(executor, executor.getFlow(), workerTaskResults);
+        }
+
+
+        if (workerTasks.isEmpty() || hasMockedWorkerTask) {
             return executor;
         }
 
diff --git a/core/src/main/java/io/kestra/core/test/Fixture.java b/core/src/main/java/io/kestra/core/test/Fixture.java
new file mode 100644
index 00000000000..387dff5ede9
--- /dev/null
+++ b/core/src/main/java/io/kestra/core/test/Fixture.java
@@ -0,0 +1,23 @@
+package io.kestra.core.test;
+
+import io.kestra.core.models.flows.State;
+import jakarta.annotation.Nullable;
+import jakarta.validation.constraints.NotNull;
+import lombok.Builder;
+import lombok.Getter;
+
+import java.util.Map;
+
+@Getter
+@Builder
+public class Fixture {
+    @NotNull
+    private String id;
+
+    private String value;
+
+    @Builder.Default
+    private State.Type state = State.Type.SUCCESS;
+
+    private Map<String, Object> outputs;
+}
diff --git a/core/src/main/java/io/kestra/core/test/Input.java b/core/src/main/java/io/kestra/core/test/Input.java
new file mode 100644
index 00000000000..1c975611484
--- /dev/null
+++ b/core/src/main/java/io/kestra/core/test/Input.java
@@ -0,0 +1,15 @@
+package io.kestra.core.test;
+
+import jakarta.validation.constraints.NotNull;
+import lombok.Builder;
+import lombok.Getter;
+
+@Getter
+@Builder
+public class Input {
+    @NotNull
+    private String id;
+
+    @NotNull
+    private String value;
+}
diff --git a/core/src/main/java/io/kestra/core/test/Test.java b/core/src/main/java/io/kestra/core/test/Test.java
new file mode 100644
index 00000000000..5466407e942
--- /dev/null
+++ b/core/src/main/java/io/kestra/core/test/Test.java
@@ -0,0 +1,24 @@
+package io.kestra.core.test;
+
+import jakarta.validation.constraints.NotNull;
+import lombok.Builder;
+import lombok.Getter;
+
+import java.util.List;
+
+@Getter
+@Builder
+public class Test {
+    @NotNull
+    private String id;
+
+    @NotNull
+    private String namespace;
+
+    @NotNull
+    private String flowId;
+
+    private List<Input> inputs;
+
+    private List<Fixture> fixtures;
+}
diff --git a/test1.yaml b/test1.yaml
new file mode 100644
index 00000000000..1987afbebfb
--- /dev/null
+++ b/test1.yaml
@@ -0,0 +1,21 @@
+id: test1
+namespace: company.team
+flowId: kookabura_411641
+
+# Define inputs
+inputs:
+  - id: input
+    value: toto
+
+# Define fixtures: a.k.a. mock of task runs and their outputs
+fixtures:
+  # Don't run it and set its output
+  - id: output
+    outputs:
+      values:
+        name: fake
+  # Don't run it
+  - id: sleep
+  # Don't run it and set its state to WARNING 
+  - id: fail
+    state: WARNING
diff --git a/test2.yaml b/test2.yaml
new file mode 100644
index 00000000000..0876ca72ba5
--- /dev/null
+++ b/test2.yaml
@@ -0,0 +1,13 @@
+id: test2
+namespace: company.team
+flowId: dogfish_450884
+
+# Define inputs
+inputs:
+  - id: input
+    value: fake
+
+# Define fixtures: a.k.a. mock of task runs and their outputs
+fixtures:
+  # Don't run it
+  - id: start
diff --git a/webserver/src/main/java/io/kestra/webserver/controllers/api/ExecutionController.java b/webserver/src/main/java/io/kestra/webserver/controllers/api/ExecutionController.java
index 4a2df932656..a034fb136a3 100644
--- a/webserver/src/main/java/io/kestra/webserver/controllers/api/ExecutionController.java
+++ b/webserver/src/main/java/io/kestra/webserver/controllers/api/ExecutionController.java
@@ -33,6 +33,7 @@
 import io.kestra.core.storages.StorageContext;
 import io.kestra.core.storages.StorageInterface;
 import io.kestra.core.tenant.TenantService;
+import io.kestra.core.test.Fixture;
 import io.kestra.core.trace.propagation.ExecutionTextMapSetter;
 import io.kestra.core.utils.Await;
 import io.kestra.core.utils.ListUtils;
@@ -694,8 +695,8 @@ public static class ExecutionResponse extends Execution {
         private final URI url;
 
         // This is not nice, but we cannot use @AllArgsConstructor as it would open a bunch of necessary changes on the Execution class.
-        ExecutionResponse(String tenantId, String id, String namespace, String flowId, Integer flowRevision, List<TaskRun> taskRunList, Map<String, Object> inputs, Map<String, Object> outputs, List<Label> labels, Map<String, Object> variables, State state, String parentId, String originalId, ExecutionTrigger trigger, boolean deleted, ExecutionMetadata metadata, Instant scheduleDate, String traceParent, URI url) {
-            super(tenantId, id, namespace, flowId, flowRevision, taskRunList, inputs, outputs, labels, variables, state, parentId, originalId, trigger, deleted, metadata, scheduleDate, traceParent);
+        ExecutionResponse(String tenantId, String id, String namespace, String flowId, Integer flowRevision, List<TaskRun> taskRunList, Map<String, Object> inputs, Map<String, Object> outputs, List<Label> labels, Map<String, Object> variables, State state, String parentId, String originalId, ExecutionTrigger trigger, boolean deleted, ExecutionMetadata metadata, Instant scheduleDate, String traceParent, List<Fixture> fixtures, URI url) {
+            super(tenantId, id, namespace, flowId, flowRevision, taskRunList, inputs, outputs, labels, variables, state, parentId, originalId, trigger, deleted, metadata, scheduleDate, traceParent, fixtures);
 
             this.url = url;
         }
@@ -720,6 +721,7 @@ public static ExecutionResponse fromExecution(Execution execution, URI url) {
                 execution.getMetadata(),
                 execution.getScheduleDate(),
                 execution.getTraceParent(),
+                execution.getFixtures(),
                 url
             );
         }
diff --git a/webserver/src/main/java/io/kestra/webserver/controllers/api/TestController.java b/webserver/src/main/java/io/kestra/webserver/controllers/api/TestController.java
new file mode 100644
index 00000000000..1091efe148e
--- /dev/null
+++ b/webserver/src/main/java/io/kestra/webserver/controllers/api/TestController.java
@@ -0,0 +1,100 @@
+package io.kestra.webserver.controllers.api;
+
+import io.kestra.core.events.CrudEvent;
+import io.kestra.core.events.CrudEventType;
+import io.kestra.core.models.Label;
+import io.kestra.core.models.executions.Execution;
+import io.kestra.core.models.executions.TaskRun;
+import io.kestra.core.models.flows.Flow;
+import io.kestra.core.models.flows.State;
+import io.kestra.core.queues.QueueException;
+import io.kestra.core.queues.QueueFactoryInterface;
+import io.kestra.core.queues.QueueInterface;
+import io.kestra.core.serializers.YamlParser;
+import io.kestra.core.services.FlowService;
+import io.kestra.core.tenant.TenantService;
+import io.kestra.core.test.Test;
+import io.kestra.core.utils.IdUtils;
+import io.kestra.core.utils.ListUtils;
+import io.micronaut.context.event.ApplicationEventPublisher;
+import io.micronaut.http.MediaType;
+import io.micronaut.http.annotation.Body;
+import io.micronaut.http.annotation.Controller;
+import io.micronaut.http.annotation.Post;
+import io.micronaut.scheduling.TaskExecutors;
+import io.micronaut.scheduling.annotation.ExecuteOn;
+import io.micronaut.validation.Validated;
+import io.swagger.v3.oas.annotations.Operation;
+import io.swagger.v3.oas.annotations.Parameter;
+import jakarta.inject.Inject;
+import jakarta.inject.Named;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+import static io.kestra.core.utils.Rethrow.throwFunction;
+
+@Slf4j
+@Validated
+@Controller("/api/v1/tests")
+public class TestController {
+    @Inject
+    private FlowService flowService;
+
+    @Inject
+    private TenantService tenantService;
+
+    @Inject
+    private YamlParser yamlParser;
+
+    @Inject
+    @Named(QueueFactoryInterface.EXECUTION_NAMED)
+    protected QueueInterface<Execution> executionQueue;
+
+    @Inject
+    private ApplicationEventPublisher<CrudEvent<Execution>> eventPublisher;
+
+    @ExecuteOn(TaskExecutors.IO)
+    @Post(consumes = MediaType.APPLICATION_YAML)
+    @Operation(tags = {"Flows"}, summary = "Test a flow")
+    public void test(@Parameter(description = "The flow namespace") @Body String testStr) throws QueueException, InterruptedException {
+        Test test = yamlParser.parse(testStr, Test.class);
+
+        Flow flow = flowService.getFlowIfExecutableOrThrow(tenantService.resolveTenant(), test.getNamespace(), test.getFlowId(), Optional.empty());
+        Map<String, Object> inputs = ListUtils.emptyOnNull(test.getInputs()).stream().collect(Collectors.toMap(
+            input -> input.getId(),
+            input -> input.getValue()
+        ));
+        Execution current = Execution.newExecution(flow, (f, e) -> inputs, null, Optional.empty())
+            .withFixtures(test.getFixtures());
+
+
+        List<Label> existingLabels = ListUtils.emptyOnNull(current.getLabels());
+        existingLabels.add(new Label(Label.TEST, "true"));
+        current = current.withLabels(existingLabels);
+
+//        List<TaskRun> taskRuns = ListUtils.emptyOnNull(test.getFixtures()).stream()
+//            .map(throwFunction(fixture -> {
+//                Thread.sleep(1); // This is to have correct ordering based on task run start date
+//                return TaskRun.builder()
+//                        .id(IdUtils.create())
+//                        .state(new State().withState(Optional.ofNullable(fixture.getState()).orElse(State.Type.SUCCESS)))
+//                        .taskId(fixture.getId())
+//                        .tenantId(flow.getTenantId())
+//                        .namespace(test.getNamespace())
+//                        .flowId(test.getFlowId())
+//                        .outputs(fixture.getOutputs())
+//                        .build();
+//                }
+//            ))
+//            .toList();
+//        current = current.withTaskRunList(taskRuns);
+
+        executionQueue.emit(current);
+        eventPublisher.publishEvent(new CrudEvent<>(current, CrudEventType.CREATE));
+    }
+
+}