add chaining to flow.Connect

Author: ezynda3

flyt.go

@@ -728,6 +728,10 @@ func Run(ctx context.Context, node Node, shared *SharedStore) (Action, error) {
 //	flow.Connect(validateNode, ActionSuccess, processNode)
 //	flow.Connect(validateNode, ActionFail, errorNode)
 //
+//	// Or with chaining:
+//	flow.Connect(validateNode, ActionSuccess, processNode)
+//	     .Connect(validateNode, ActionFail, errorNode)
+//
 //	// Run flow
 //	err := flow.Run(ctx, shared)
 type Flow struct {
@@ -757,6 +761,7 @@ func NewFlow(start Node) *Flow {
 // Connect adds a transition from one node to another based on an action.
 // When the 'from' node returns the specified action, execution continues
 // with the 'to' node. Multiple actions can be connected from a single node.
+// Returns the flow instance for method chaining.
 //
 // Parameters:
 //   - from: The source node
@@ -768,11 +773,18 @@ func NewFlow(start Node) *Flow {
 //	flow.Connect(nodeA, "success", nodeB)
 //	flow.Connect(nodeA, "retry", nodeA)  // Self-loop for retry
 //	flow.Connect(nodeA, "fail", errorNode)
-func (f *Flow) Connect(from Node, action Action, to Node) {
+//
+// Example with chaining:
+//
+//	flow.Connect(nodeA, "success", nodeB)
+//	     .Connect(nodeB, "success", finalNode)
+//	     .Connect(nodeB, "fail", errorNode)
+func (f *Flow) Connect(from Node, action Action, to Node) *Flow {
 	if f.transitions[from] == nil {
 		f.transitions[from] = make(map[Action]Node)
 	}
 	f.transitions[from][action] = to
+	return f
 }
 
 // Run executes the flow starting from the start node.

flyt_test.go

@@ -689,4 +689,142 @@ func BenchmarkSharedStoreConcurrentAccess(b *testing.B) {
 	})
 }
 
+// TestFlowConnectChaining tests that Connect method can be chained
+func TestFlowConnectChaining(t *testing.T) {
+	// Create nodes
+	node1 := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "node1",
+		action:   "next",
+	}
+
+	node2 := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "node2",
+		action:   "final",
+	}
+
+	node3 := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "node3",
+		action:   DefaultAction,
+	}
+
+	// Test chaining
+	flow := NewFlow(node1)
+	flow.Connect(node1, "next", node2).
+		Connect(node2, "final", node3)
+
+	// Run the flow
+	ctx := context.Background()
+	shared := NewSharedStore()
+	err := flow.Run(ctx, shared)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// Verify execution order
+	if !node1.executed {
+		t.Error("node1 not executed")
+	}
+	if !node2.executed {
+		t.Error("node2 not executed")
+	}
+	if !node3.executed {
+		t.Error("node3 not executed")
+	}
+}
+
+// TestFlowConnectChainingWithMultipleActions tests chaining with multiple actions from same node
+func TestFlowConnectChainingWithMultipleActions(t *testing.T) {
+	// Create nodes
+	decisionNode := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "decision",
+		action:   "branch",
+	}
+
+	successNode := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "success",
+		action:   DefaultAction,
+	}
+
+	failNode := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "fail",
+		action:   DefaultAction,
+	}
+
+	retryNode := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "retry",
+		action:   DefaultAction,
+	}
+
+	// Test chaining with multiple actions from same node
+	flow := NewFlow(decisionNode)
+	flow.Connect(decisionNode, "branch", successNode).
+		Connect(decisionNode, "fail", failNode).
+		Connect(decisionNode, "retry", retryNode)
+
+	// Run the flow
+	ctx := context.Background()
+	shared := NewSharedStore()
+	err := flow.Run(ctx, shared)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// Verify execution
+	if !decisionNode.executed {
+		t.Error("decisionNode not executed")
+	}
+	if !successNode.executed {
+		t.Error("successNode not executed")
+	}
+	if failNode.executed {
+		t.Error("failNode should not be executed")
+	}
+	if retryNode.executed {
+		t.Error("retryNode should not be executed")
+	}
+}
+
+// TestFlowConnectChainingBackwardsCompatibility tests that traditional usage still works
+func TestFlowConnectChainingBackwardsCompatibility(t *testing.T) {
+	// Create nodes
+	node1 := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "node1",
+		action:   "next",
+	}
+
+	node2 := &TestNode{
+		BaseNode: NewBaseNode(),
+		name:     "node2",
+		action:   DefaultAction,
+	}
+
+	// Test that traditional non-chaining usage still works
+	flow := NewFlow(node1)
+	flow.Connect(node1, "next", node2) // Traditional usage without chaining
+
+	// Run the flow
+	ctx := context.Background()
+	shared := NewSharedStore()
+	err := flow.Run(ctx, shared)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// Verify execution
+	if !node1.executed {
+		t.Error("node1 not executed")
+	}
+	if !node2.executed {
+		t.Error("node2 not executed")
+	}
+}
+
 // BenchmarkBatchNodeProcessing benchmarks batch processing