Fix: Correctly set 'child data dir/output dir' in branch handler's Abort and Finalize

flytepropeller/pkg/controller/nodes/branch/handler.go

@@ -123,6 +123,25 @@ func (b *branchHandler) getExecutionContextForDownstream(nCtx interfaces.NodeExe
 	return executors.NewExecutionContextWithParentInfo(nCtx.ExecutionContext(), newParentInfo), nil
 }
 
+// setupBranchTakenNodePaths sets the DataDir and OutputDir on the branch-taken node's
+// status using the branch node's OutputDir as the base, and returns the computed OutputDir.
+// This must be called before recursing into the branch-taken node so that any nested handler
+// (e.g. a dynamic node) finds the correct paths.
+func (b *branchHandler) setupBranchTakenNodePaths(ctx context.Context, nCtx interfaces.NodeExecutionContext, branchTakenNode v1alpha1.ExecutableNode) (storage.DataReference, error) {
+	childNodeStatus := nCtx.ContextualNodeLookup().GetNodeExecutionStatus(ctx, branchTakenNode.GetID())
+	childDataDir, err := nCtx.DataStore().ConstructReference(ctx, nCtx.NodeStatus().GetOutputDir(), branchTakenNode.GetID())
+	if err != nil {
+		return "", err
+	}
+	childOutputDir, err := nCtx.DataStore().ConstructReference(ctx, childDataDir, strconv.Itoa(int(childNodeStatus.GetAttempts())))
+	if err != nil {
+		return "", err
+	}
+	childNodeStatus.SetDataDir(childDataDir)
+	childNodeStatus.SetOutputDir(childOutputDir)
+	return childOutputDir, nil
+}
+
 func (b *branchHandler) recurseDownstream(ctx context.Context, nCtx interfaces.NodeExecutionContext, branchTakenNode v1alpha1.ExecutableNode) (handler.Transition, error) {
 	// TODO we should replace the call to RecursiveNodeHandler with a call to SingleNode Handler. The inputs are also already known ahead of time
 	// There is no DAGStructure for the branch nodes, the branch taken node is the leaf node. The node itself may be arbitrarily complex, but in that case the node should reference a subworkflow etc
@@ -133,17 +152,10 @@ func (b *branchHandler) recurseDownstream(ctx context.Context, nCtx interfaces.N
 			errors.Errorf(errors.IllegalStateError, nCtx.NodeID(), "nodeLookup must be supplied.")
 	}
 
-	childNodeStatus := nl.GetNodeExecutionStatus(ctx, branchTakenNode.GetID())
-	childDataDir, err := nCtx.DataStore().ConstructReference(ctx, nCtx.NodeStatus().GetOutputDir(), branchTakenNode.GetID())
+	childOutputDir, err := b.setupBranchTakenNodePaths(ctx, nCtx, branchTakenNode)
 	if err != nil {
 		return handler.UnknownTransition, err
 	}
-	childOutputDir, err := nCtx.DataStore().ConstructReference(ctx, childDataDir, strconv.Itoa(int(childNodeStatus.GetAttempts())))
-	if err != nil {
-		return handler.UnknownTransition, err
-	}
-	childNodeStatus.SetDataDir(childDataDir)
-	childNodeStatus.SetOutputDir(childOutputDir)
 	upstreamNodeIds, err := nCtx.ContextualNodeLookup().ToNode(branchTakenNode.GetID())
 	if err != nil {
 		return handler.UnknownTransition, err
@@ -213,6 +225,9 @@ func (b *branchHandler) Abort(ctx context.Context, nCtx interfaces.NodeExecution
 	// TODO we should replace the call to RecursiveNodeHandler with a call to SingleNode Handler. The inputs are also already known ahead of time
 	// There is no DAGStructure for the branch nodes, the branch taken node is the leaf node. The node itself may be arbitrarily complex, but in that case the node should reference a subworkflow etc
 	// The parent of the BranchTaken Node is the actual Branch Node and all the data is just forwarded from the Branch to the executed node.
+	if _, err := b.setupBranchTakenNodePaths(ctx, nCtx, branchTakenNode); err != nil {
+		return err
+	}
 	upstreamNodeIds, err := nCtx.ContextualNodeLookup().ToNode(branchTakenNode.GetID())
 	if err != nil {
 		return err
@@ -257,6 +272,9 @@ func (b *branchHandler) Finalize(ctx context.Context, nCtx interfaces.NodeExecut
 	// TODO we should replace the call to RecursiveNodeHandler with a call to SingleNode Handler. The inputs are also already known ahead of time
 	// There is no DAGStructure for the branch nodes, the branch taken node is the leaf node. The node itself may be arbitrarily complex, but in that case the node should reference a subworkflow etc
 	// The parent of the BranchTaken Node is the actual Branch Node and all the data is just forwarded from the Branch to the executed node.
+	if _, err := b.setupBranchTakenNodePaths(ctx, nCtx, branchTakenNode); err != nil {
+		return err
+	}
 	upstreamNodeIds, err := nCtx.ContextualNodeLookup().ToNode(branchTakenNode.GetID())
 	if err != nil {
 		return err

flytepropeller/pkg/controller/nodes/branch/handler_test.go

@@ -336,9 +336,75 @@ func TestBranchHandler_AbortNode(t *testing.T) {
 			mock.Anything,
 			mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		mockNodeLookup.EXPECT().GetNode(*s.s.FinalizedNodeID).Return(n, true)
+		childNodeStatus := &mocks2.ExecutableNodeStatus{}
+		childNodeStatus.EXPECT().GetAttempts().Return(0)
+		childNodeStatus.On("SetDataDir", storage.DataReference("/output-dir/n2")).Once()
+		childNodeStatus.On("SetOutputDir", storage.DataReference("/output-dir/n2/0")).Once()
+		mockNodeLookup.EXPECT().GetNodeExecutionStatus(ctx, n2).Return(childNodeStatus)
 		branch := New(mockNodeExecutor, eventConfig, promutils.NewTestScope())
 		err := branch.Abort(ctx, nCtx, "")
 		assert.NoError(t, err)
+		childNodeStatus.AssertExpectations(t)
+	})
+}
+
+func TestBranchHandler_FinalizeNode(t *testing.T) {
+	ctx := context.TODO()
+	n1 := "n1"
+	n2 := "n2"
+
+	exp, _ := getComparisonExpression(1.0, core.ComparisonExpression_EQ, 1.0)
+	branchNode := &v1alpha1.BranchNodeSpec{
+		If: v1alpha1.IfBlock{
+			Condition: v1alpha1.BooleanExpression{
+				BooleanExpression: &core.BooleanExpression{
+					Expr: &core.BooleanExpression_Comparison{
+						Comparison: exp,
+					},
+				},
+			},
+			ThenNode: &n1,
+		},
+	}
+
+	n := &v1alpha1.NodeSpec{
+		ID:         n2,
+		BranchNode: branchNode,
+	}
+
+	t.Run("NoBranchNode", func(t *testing.T) {
+		mockNodeExecutor := &mocks.Node{}
+		eCtx := &execMocks.ExecutionContext{}
+		eCtx.EXPECT().GetParentInfo().Return(nil)
+		nCtx, _ := createNodeContext(v1alpha1.BranchNodeError, nil, n, nil, nil, eCtx)
+		branch := New(mockNodeExecutor, eventConfig, promutils.NewTestScope())
+		err := branch.Finalize(ctx, nCtx)
+		assert.NoError(t, err)
+	})
+
+	t.Run("BranchNodeSuccess", func(t *testing.T) {
+		mockNodeExecutor := &mocks.Node{}
+		mockNodeLookup := &execMocks.NodeLookup{}
+		mockNodeLookup.EXPECT().ToNode(mock.Anything).Return(nil, nil)
+		eCtx := &execMocks.ExecutionContext{}
+		eCtx.EXPECT().GetParentInfo().Return(parentInfo{})
+		nCtx, s := createNodeContext(v1alpha1.BranchNodeSuccess, &n1, n, nil, mockNodeLookup, eCtx)
+		newParentInfo, _ := common.CreateParentInfo(parentInfo{}, nCtx.NodeID(), nCtx.CurrentAttempt(), false)
+		expectedExecContext := executors.NewExecutionContextWithParentInfo(nCtx.ExecutionContext(), newParentInfo)
+		mockNodeExecutor.EXPECT().FinalizeHandler(mock.Anything,
+			mock.MatchedBy(func(e executors.ExecutionContext) bool { return assert.Equal(t, e, expectedExecContext) }),
+			mock.Anything,
+			mock.Anything, mock.Anything).Return(nil)
+		mockNodeLookup.EXPECT().GetNode(*s.s.FinalizedNodeID).Return(n, true)
+		childNodeStatus := &mocks2.ExecutableNodeStatus{}
+		childNodeStatus.EXPECT().GetAttempts().Return(0)
+		childNodeStatus.On("SetDataDir", storage.DataReference("/output-dir/n2")).Once()
+		childNodeStatus.On("SetOutputDir", storage.DataReference("/output-dir/n2/0")).Once()
+		mockNodeLookup.EXPECT().GetNodeExecutionStatus(ctx, n2).Return(childNodeStatus)
+		branch := New(mockNodeExecutor, eventConfig, promutils.NewTestScope())
+		err := branch.Finalize(ctx, nCtx)
+		assert.NoError(t, err)
+		childNodeStatus.AssertExpectations(t)
 	})
 }