Explicitly fall back to real test execution when testing (#218)

Explicitly fall back to real test execution when testing

Instead of silently falling back to real task execution when mocked inputs are unmatched, this now requires an explicit mock passing in the task's `run` method.

Signed-off-by: Hongxin Liang <[email protected]>
Co-authored-by: Hongxin Liang <[email protected]>

Author: honnix

flytekit-testing/src/main/java/org/flyte/flytekit/testing/TestingRunnableLaunchPlan.java

@@ -35,12 +35,14 @@ public class TestingRunnableLaunchPlan<InputT, OutputT>
       SdkType<InputT> inputType,
       SdkType<OutputT> outputType,
       Function<InputT, OutputT> runFn,
+      boolean runFnProvided,
       Map<InputT, OutputT> fixedOutputs) {
     super(
         launchPlanId,
         inputType,
         outputType,
         runFn,
+        runFnProvided,
         fixedOutputs,
         TestingRunnableLaunchPlan::new,
         "launch plan",
@@ -52,6 +54,7 @@ static <InputT, OutputT> TestingRunnableLaunchPlan<InputT, OutputT> create(
     PartialLaunchPlanIdentifier launchPlanId =
         PartialLaunchPlanIdentifier.builder().name(name).build();
 
-    return new TestingRunnableLaunchPlan<>(launchPlanId, inputType, outputType, null, emptyMap());
+    return new TestingRunnableLaunchPlan<>(
+        launchPlanId, inputType, outputType, null, false, emptyMap());
   }
 }

flytekit-testing/src/main/java/org/flyte/flytekit/testing/TestingRunnableNode.java

@@ -38,6 +38,7 @@ public abstract class TestingRunnableNode<
 
   // @Nullable - signal nullable field but without adding the dependency
   protected final Function<InputT, OutputT> runFn;
+  private final boolean runFnProvided;
 
   protected final Map<InputT, OutputT> fixedOutputs;
   private final Creator<IdT, InputT, OutputT, T> creatorFn;
@@ -54,6 +55,7 @@ T create(
         SdkType<InputT> inputType,
         SdkType<OutputT> outputType,
         Function<InputT, OutputT> runFn,
+        boolean runFnProvided,
         Map<InputT, OutputT> fixedOutputs);
   }
 
@@ -62,6 +64,7 @@ protected TestingRunnableNode(
       SdkType<InputT> inputType,
       SdkType<OutputT> outputType,
       Function<InputT, OutputT> runFn,
+      boolean runFnProvided,
       Map<InputT, OutputT> fixedOutputs,
       Creator<IdT, InputT, OutputT, T> creatorFn,
       String type,
@@ -70,6 +73,7 @@ protected TestingRunnableNode(
     this.inputType = requireNonNull(inputType, "inputType");
     this.outputType = requireNonNull(outputType, "outputType");
     this.runFn = runFn; // Nullable
+    this.runFnProvided = runFnProvided;
     this.fixedOutputs = requireNonNull(fixedOutputs, "fixedOutputs");
     this.creatorFn = requireNonNull(creatorFn, "creatorFn");
     this.type = requireNonNull(type, "type");
@@ -80,19 +84,28 @@ protected TestingRunnableNode(
   public Map<String, Literal> run(Map<String, Literal> inputs) {
     InputT input = inputType.fromLiteralMap(inputs);
 
-    if (fixedOutputs.containsKey(input)) {
-      return outputType.toLiteralMap(fixedOutputs.get(input));
-    } else if (runFn != null) {
-      return outputType.toLiteralMap(runFn.apply(input));
+    if (fixedOutputs.size() == 0) {
+      // No mocking via input matching, either run the real thing or run the provided lambda
+      if (runFn != null) {
+        return outputType.toLiteralMap(runFn.apply(input));
+      }
+    } else {
+      if (fixedOutputs.containsKey(input)) {
+        return outputType.toLiteralMap(fixedOutputs.get(input));
+      }
+      // Inputs not matching, run the provided lambda
+      if (runFn != null && runFnProvided) {
+        return outputType.toLiteralMap(runFn.apply(input));
+      }
     }
 
-    // TODO see if we can improve this error message as input is hard to read
-    // We can improve the SdkBindingData toString method
     String message =
         String.format(
             "Can't find input %s for remote %s [%s] across known %s inputs, "
                 + "use %s to provide a test double",
             input, type, getName(), type, testingSuggestion);
+
+    // Not matching inputs and there is nothing to run
     throw new IllegalArgumentException(message);
   }
 
@@ -105,10 +118,10 @@ public T withFixedOutput(InputT input, OutputT output) {
     Map<InputT, OutputT> newFixedOutputs = new HashMap<>(fixedOutputs);
     newFixedOutputs.put(input, output);
 
-    return creatorFn.create(id, inputType, outputType, runFn, newFixedOutputs);
+    return creatorFn.create(id, inputType, outputType, runFn, runFnProvided, newFixedOutputs);
   }
 
   public T withRunFn(Function<InputT, OutputT> runFn) {
-    return creatorFn.create(id, inputType, outputType, runFn, fixedOutputs);
+    return creatorFn.create(id, inputType, outputType, runFn, true, fixedOutputs);
   }
 }

flytekit-testing/src/main/java/org/flyte/flytekit/testing/TestingRunnableTask.java

@@ -37,12 +37,14 @@ private TestingRunnableTask(
       SdkType<InputT> inputType,
       SdkType<OutputT> outputType,
       Function<InputT, OutputT> runFn,
+      boolean runFnProvided,
       Map<InputT, OutputT> fixedOutputs) {
     super(
         taskId,
         inputType,
         outputType,
         runFn,
+        runFnProvided,
         fixedOutputs,
         TestingRunnableTask::new,
         "task",
@@ -54,14 +56,15 @@ static <InputT, OutputT> TestingRunnableTask<InputT, OutputT> create(
     PartialTaskIdentifier taskId = PartialTaskIdentifier.builder().name(task.getName()).build();
 
     return new TestingRunnableTask<>(
-        taskId, task.getInputType(), task.getOutputType(), task::run, emptyMap());
+        taskId, task.getInputType(), task.getOutputType(), task::run, false, emptyMap());
   }
 
   static <InputT, OutputT> TestingRunnableTask<InputT, OutputT> create(
       String name, SdkType<InputT> inputType, SdkType<OutputT> outputType) {
     PartialTaskIdentifier taskId = PartialTaskIdentifier.builder().name(name).build();
 
-    return new TestingRunnableTask<>(taskId, inputType, outputType, /* runFn= */ null, emptyMap());
+    return new TestingRunnableTask<>(
+        taskId, inputType, outputType, /* runFn= */ null, false, emptyMap());
   }
 
   @Override

flytekit-testing/src/test/java/org/flyte/flytekit/testing/FibonacciWorkflowTest.java

@@ -73,6 +73,7 @@ public void testWithTaskOutput_runnableTask() {
                 new SumTask(),
                 SumInput.create(SdkBindingDataFactory.of(3L), SdkBindingDataFactory.of(5L)),
                 SumOutput.create(SdkBindingDataFactory.of(42L)))
+            .withTask(new SumTask(), new SumTask()::run)
             .execute();
 
     assertThat(result.getIntegerOutput("fib2"), equalTo(2L));

flytekit-testing/src/test/java/org/flyte/flytekit/testing/TestingRunnableNodeTest.java

@@ -85,6 +85,17 @@ void testRun_notFound() {
                 + "across known test inputs, use a magic wang to provide a test double"));
   }
 
+  @Test
+  void testRun_notFoundRunFnProvided() {
+    Function<Input, Output> fn = in -> Output.create(Long.parseLong(in.in().get()));
+    Map<Input, Output> fixedOutputs = singletonMap(Input.create("7"), Output.create(7L));
+    TestNode node = new TestNode(fn, fixedOutputs);
+
+    Map<String, Literal> output = node.run(singletonMap("in", Literals.ofString("10")));
+
+    assertThat(output, hasEntry("out", Literals.ofInteger(10L)));
+  }
+
   @Test
   void testWithFixedOutput() {
     TestNode node =
@@ -105,6 +116,22 @@ void testWithRunFn() {
     assertThat(output, hasEntry("out", Literals.ofInteger(7L)));
   }
 
+  @Test
+  void testWithoutRunFn() {
+    TestNode node = new TestNode(null, emptyMap());
+
+    IllegalArgumentException ex =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> node.run(singletonMap("in", Literals.ofString("7"))));
+
+    assertThat(
+        ex.getMessage(),
+        equalTo(
+            "Can't find input Input{in=SdkBindingData{type=strings, value=7}} for remote test [TestTask] "
+                + "across known test inputs, use a magic wang to provide a test double"));
+  }
+
   @Test
   void testGetNameShouldDeriveFromId() {
     TestNode node = new TestNode(null, emptyMap());
@@ -121,8 +148,9 @@ protected TestNode(Function<Input, Output> runFn, Map<Input, Output> fixedOutput
           JacksonSdkType.of(Input.class),
           JacksonSdkType.of(Output.class),
           runFn,
+          true,
           fixedOutputs,
-          (id, inType, outType, f, m) -> new TestNode(f, m),
+          (id, inType, outType, f, fProvided, m) -> new TestNode(f, m),
           "test",
           "a magic wang");
     }