feat: Add start node for DOT format

Signed-off-by: Sergiy Kulanov <[email protected]>
SergK · Oct 8, 2023 · 74e7618 · 74e7618
1 parent 2039799
commit 74e7618
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/pkg/taskgraph/taskgraph.go b/pkg/taskgraph/taskgraph.go
@@ -74,6 +74,10 @@ func (g *TaskGraph) ToDOT() *DOT {
 	}
 
 	for _, node := range g.Nodes {
+		if !node.hasParent {
+			// "start" is the special node that represents the start of the pipeline
+			dot.Edges = append(dot.Edges, fmt.Sprintf("  \"start\" -> \"%s\"", node.Name))
+		}
 		if len(node.Dependencies) == 0 {
 			// "end" is the special node that represents the end of the pipeline
 			dot.Edges = append(dot.Edges, fmt.Sprintf("  \"%s\" -> \"end\"", node.Name))
@@ -94,6 +98,10 @@ func (g *TaskGraph) ToDOTWithTaskRef() *DOT {
 	}
 
 	for _, node := range g.Nodes {
+		if !node.hasParent {
+			// "start" is the special node that represents the start of the pipeline
+			dot.Edges = append(dot.Edges, fmt.Sprintf("  \"start\" -> \"%s\n(%s)\"", node.Name, node.TaskRefName))
+		}
 		if len(node.Dependencies) == 0 {
 			// "end" is the special node that represents the end of the pipeline
 			dot.Edges = append(dot.Edges, fmt.Sprintf("  \"%s\n(%s)\" -> \"end\"", node.Name, node.TaskRefName))
@@ -109,7 +117,7 @@ func (g *TaskGraph) ToDOTWithTaskRef() *DOT {
 // String converts a DOT graph to a string
 func (d *DOT) String() string {
 	var buf bytes.Buffer
-	buf.WriteString(fmt.Sprintf("%s {\n  labelloc=\"t\"\n  label=\"%s\"\n  end [shape=\"point\" width=0.2]\n", d.Format, d.Name))
+	buf.WriteString(fmt.Sprintf("%s {\n  labelloc=\"t\"\n  label=\"%s\"\n  end [shape=\"point\" width=0.2]\n  start [shape=\"point\" width=0.2]\n", d.Format, d.Name))
 	for _, edge := range d.Edges {
 		buf.WriteString(fmt.Sprintf("%s\n", edge))
 	}

diff --git a/pkg/taskgraph/taskgraph_test.go b/pkg/taskgraph/taskgraph_test.go
@@ -80,6 +80,8 @@ func TestTaskGraphToDOT(t *testing.T) {
 	assert.Contains(t, dot.Edges, "  \"task3\" -> \"task2\"")
 	assert.Contains(t, dot.Edges, "  \"task1\" -> \"end\"")
 	assert.Contains(t, dot.Edges, "  \"task4\" -> \"end\"")
+	assert.Contains(t, dot.Edges, "  \"start\" -> \"task4\"")
+	assert.Contains(t, dot.Edges, "  \"start\" -> \"task3\"")
 }
 
 func TestTaskGraphToDOTWithTaskRef(t *testing.T) {
@@ -96,6 +98,8 @@ func TestTaskGraphToDOTWithTaskRef(t *testing.T) {
 	assert.Contains(t, dot.Edges, "  \"task3\n(taskRef3)\" -> \"task2\n(taskRef2)\"")
 	assert.Contains(t, dot.Edges, "  \"task1\n(taskRef1)\" -> \"end\"")
 	assert.Contains(t, dot.Edges, "  \"task4\n(taskRef4)\" -> \"end\"")
+	assert.Contains(t, dot.Edges, "  \"start\" -> \"task3\n(taskRef3)\"")
+	assert.Contains(t, dot.Edges, "  \"start\" -> \"task4\n(taskRef4)\"")
 }
 
 func TestDOTString(t *testing.T) {
@@ -114,7 +118,16 @@ func TestDOTString(t *testing.T) {
 	}
 
 	// Test the String method
-	expected := "digraph {\n  labelloc=\"t\"\n  label=\"test-pipeline\"\n  end [shape=\"point\" width=0.2]\n  \"task1\" -> \"task2\"\n  \"task1\" -> \"task3\"\n  \"task2\" -> \"task3\"\n}\n"
+	expected := `digraph {
+  labelloc="t"
+  label="test-pipeline"
+  end [shape="point" width=0.2]
+  start [shape="point" width=0.2]
+  "task1" -> "task2"
+  "task1" -> "task3"
+  "task2" -> "task3"
+}
+`
 	assert.Equal(t, expected, dot.String())
 }