include node id when missing mocks:

Signed-off-by: Anderson Reyes <[email protected]>
flyteorg · Sep 22, 2023 · a83d2b1 · a83d2b1
1 parent ea7c1aa
commit a83d2b1
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 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
@@ -219,25 +219,27 @@ 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,17 +249,18 @@ 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:
         String subWorkflowName = reference.subWorkflowRef().name();
         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
@@ -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,53 @@ 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 =