From 7ed35134b35fd681527b6c28a22e8c2d10beb662 Mon Sep 17 00:00:00 2001
From: Anderson Reyes <andersonr@spotify.com>
Date: Fri, 22 Sep 2023 19:37:19 -0400
Subject: [PATCH] include node id when missing mocks:

Signed-off-by: Anderson Reyes <reyesanderson428@gmail.com>

formatting
---
 .../flytekit/testing/SdkTestingExecutor.java  | 25 ++++++----
 .../testing/SdkTestingExecutorTest.java       | 48 ++++++++++++++++++-
 2 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/flytekit-testing/src/main/java/org/flyte/flytekit/testing/SdkTestingExecutor.java b/flytekit-testing/src/main/java/org/flyte/flytekit/testing/SdkTestingExecutor.java
index aa55f22a9..d367ca9f3 100644
--- a/flytekit-testing/src/main/java/org/flyte/flytekit/testing/SdkTestingExecutor.java
+++ b/flytekit-testing/src/main/java/org/flyte/flytekit/testing/SdkTestingExecutor.java
@@ -219,25 +219,28 @@ private void checkInputsInFixedInputs(WorkflowTemplate template) {
   private void checkTestDoublesForNodes(WorkflowTemplate template) {
     for (Node node : template.nodes()) {
       if (node.taskNode() != null) {
-        checkTestDoubleForTaskNode(node.taskNode());
+        checkTestDoubleForTaskNode(node);
       } else if (node.workflowNode() != null) {
-        checkTestDoubleForWorkflowNode(node.workflowNode());
+        checkTestDoubleForWorkflowNode(node);
       }
     }
   }
 
-  private void checkTestDoubleForTaskNode(TaskNode taskNode) {
+  private void checkTestDoubleForTaskNode(Node node) {
+    TaskNode taskNode = node.taskNode();
     String taskName = taskNode.referenceId().name();
 
     checkArgument(
         taskTestDoubles().containsKey(taskName),
-        "Can't execute remote task [%s], "
+        "Can't execute remote task [%s] for node [%s], "
             + "use SdkTestingExecutor#withTaskOutput or SdkTestingExecutor#withTask "
             + "to provide a test double",
-        taskName);
+        taskName,
+        node.id());
   }
 
-  private void checkTestDoubleForWorkflowNode(WorkflowNode workflowNode) {
+  private void checkTestDoubleForWorkflowNode(Node node) {
+    WorkflowNode workflowNode = node.workflowNode();
     Reference reference = workflowNode.reference();
     switch (reference.kind()) {
       case LAUNCH_PLAN_REF:
@@ -247,9 +250,10 @@ private void checkTestDoubleForWorkflowNode(WorkflowNode workflowNode) {
         checkArgument(
             launchPlan != null,
             "Can't execute remote launch plan "
-                + "[%s], use SdkTestingExecutor#withLaunchPlanOutput or "
+                + "[%s] for node [%s], use SdkTestingExecutor#withLaunchPlanOutput or "
                 + "SdkTestingExecutor#withLaunchPlan to provide a test double",
-            launchPlanName);
+            launchPlanName,
+            node.id());
         return;
 
       case SUB_WORKFLOW_REF:
@@ -257,7 +261,10 @@ private void checkTestDoubleForWorkflowNode(WorkflowNode workflowNode) {
         WorkflowTemplate subWorkflowTemplate = workflowTemplates().get(subWorkflowName);
 
         checkArgument(
-            subWorkflowTemplate != null, "Can't expand sub workflow [%s]", subWorkflowName);
+            subWorkflowTemplate != null,
+            "Can't expand sub workflow [%s] for node [%s]",
+            subWorkflowName,
+            node.id());
 
         checkTestDoublesForNodes(subWorkflowTemplate);
     }
diff --git a/flytekit-testing/src/test/java/org/flyte/flytekit/testing/SdkTestingExecutorTest.java b/flytekit-testing/src/test/java/org/flyte/flytekit/testing/SdkTestingExecutorTest.java
index e9b370eba..59a6ee683 100644
--- a/flytekit-testing/src/test/java/org/flyte/flytekit/testing/SdkTestingExecutorTest.java
+++ b/flytekit-testing/src/test/java/org/flyte/flytekit/testing/SdkTestingExecutorTest.java
@@ -220,7 +220,7 @@ public Void expand(SdkWorkflowBuilder builder, Void noInput) {
     assertThat(
         e.getMessage(),
         equalTo(
-            "Can't execute remote task [remote_sum_task], use SdkTestingExecutor#withTaskOutput or "
+            "Can't execute remote task [remote_sum_task] for node [sum], use SdkTestingExecutor#withTaskOutput or "
                 + "SdkTestingExecutor#withTask to provide a test double"));
   }
 
@@ -429,6 +429,52 @@ public TestUnaryIntegerIO expand(SdkWorkflowBuilder builder, SumLaunchPlanInput
     assertThat(result.getIntegerOutput("integer"), equalTo(35L));
   }
 
+  @Test
+  public void testWithLaunchPlan_isMissingLaunchPlan() {
+    SdkRemoteLaunchPlan<SumLaunchPlanInput, SumLaunchPlanOutput> launchplanRef =
+        SdkRemoteLaunchPlan.create(
+            "development",
+            "flyte-warehouse",
+            "SumWorkflow",
+            JacksonSdkType.of(SumLaunchPlanInput.class),
+            JacksonSdkType.of(SumLaunchPlanOutput.class));
+
+    SdkWorkflow<SumLaunchPlanInput, TestUnaryIntegerIO> workflow =
+        new SdkWorkflow<>(
+            JacksonSdkType.of(SumLaunchPlanInput.class),
+            JacksonSdkType.of(TestUnaryIntegerIO.class)) {
+          @Override
+          public TestUnaryIntegerIO expand(SdkWorkflowBuilder builder, SumLaunchPlanInput input) {
+            SdkBindingData<Long> c =
+                builder
+                    .apply(
+                        "mylaunchplan",
+                        launchplanRef,
+                        SumLaunchPlanInput.create(input.a(), input.b()))
+                    .getOutputs()
+                    .c();
+
+            return TestUnaryIntegerIO.create(c);
+          }
+        };
+
+    IllegalArgumentException e =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                SdkTestingExecutor.of(workflow)
+                    .withFixedInput("a", 30L)
+                    .withFixedInput("b", 5L)
+                    .execute());
+
+    assertThat(
+        e.getMessage(),
+        equalTo(
+            "Can't execute remote launch plan "
+                + "[SumWorkflow] for node [mylaunchplan], use SdkTestingExecutor#withLaunchPlanOutput or "
+                + "SdkTestingExecutor#withLaunchPlan to provide a test double"));
+  }
+
   @Test
   public void testWithLaunchPlan_missingRemoteTaskOutput() {
     SdkWorkflow<Void, Void> workflow =