From 2039799f720464ec89932673d2163e0e9e021424 Mon Sep 17 00:00:00 2001
From: Sergiy Kulanov <sergiy_kulanov@epam.com>
Date: Sun, 8 Oct 2023 17:19:13 +0300
Subject: [PATCH] feat: Add support for start node using Mermaid format

Signed-off-by: Sergiy Kulanov <sergiy_kulanov@epam.com>
---
 pkg/taskgraph/taskgraph.go      | 13 +++++++++++--
 pkg/taskgraph/taskgraph_test.go | 12 ++++++++----
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/pkg/taskgraph/taskgraph.go b/pkg/taskgraph/taskgraph.go
index d49600f..8491222 100644
--- a/pkg/taskgraph/taskgraph.go
+++ b/pkg/taskgraph/taskgraph.go
@@ -19,6 +19,7 @@ type TaskNode struct {
 	Name         string
 	TaskRefName  string // Name of the kind: Task referenced by this task in the pipeline
 	Dependencies []*TaskNode
+	hasParent    bool // Flag to indicate if the node has a parent
 }
 
 type DOT struct {
@@ -44,6 +45,7 @@ func BuildTaskGraph(tasks []v1pipeline.PipelineTask) *TaskGraph {
 		node := &TaskNode{
 			Name:        task.Name,
 			TaskRefName: task.TaskRef.Name,
+			hasParent:   false, // we assume that the node has no parent until we find a dependency
 		}
 		graph.Nodes[task.Name] = node
 	}
@@ -57,6 +59,7 @@ func BuildTaskGraph(tasks []v1pipeline.PipelineTask) *TaskGraph {
 		for _, depName := range task.RunAfter {
 			depNode := graph.Nodes[depName]
 			depNode.Dependencies = append(depNode.Dependencies, node)
+			node.hasParent = true
 		}
 	}
 
@@ -172,7 +175,10 @@ func (g *TaskGraph) ToMermaid() string {
 	mermaid := fmt.Sprintf("---\ntitle: %s\n---\nflowchart TD\n", g.PipelineName)
 	for _, node := range g.Nodes {
 		if len(node.Dependencies) == 0 {
-			mermaid += fmt.Sprintf("   %s --> id([fa:fa-circle])\n", node.Name)
+			mermaid += fmt.Sprintf("   %s --> stop([fa:fa-circle])\n", node.Name)
+		}
+		if !node.hasParent {
+			mermaid += fmt.Sprintf("   start([fa:fa-circle]) --> %s\n", node.Name)
 		}
 		for _, dep := range node.Dependencies {
 			mermaid += fmt.Sprintf("   %s --> %s\n", node.Name, dep.Name)
@@ -186,7 +192,10 @@ func (g *TaskGraph) ToMermaidWithTaskRef() string {
 	mermaid := fmt.Sprintf("---\ntitle: %s\n---\nflowchart TD\n", g.PipelineName)
 	for _, node := range g.Nodes {
 		if len(node.Dependencies) == 0 {
-			mermaid += fmt.Sprintf("   %s(\"%s\n   (%s)\") --> id([fa:fa-circle])\n", node.Name, node.Name, node.TaskRefName)
+			mermaid += fmt.Sprintf("   %s(\"%s\n   (%s)\") --> stop([fa:fa-circle])\n", node.Name, node.Name, node.TaskRefName)
+		}
+		if !node.hasParent {
+			mermaid += fmt.Sprintf("   start([fa:fa-circle]) --> %s(\"%s\n   (%s)\")\n", node.Name, node.Name, node.TaskRefName)
 		}
 		for _, dep := range node.Dependencies {
 			mermaid += fmt.Sprintf("   %s(\"%s\n   (%s)\") --> %s(\"%s\n   (%s)\")\n", node.Name, node.Name, node.TaskRefName, dep.Name, dep.Name, dep.TaskRefName)
diff --git a/pkg/taskgraph/taskgraph_test.go b/pkg/taskgraph/taskgraph_test.go
index b1d144b..86ab69c 100644
--- a/pkg/taskgraph/taskgraph_test.go
+++ b/pkg/taskgraph/taskgraph_test.go
@@ -164,8 +164,10 @@ func TestTaskGraphToMermaid(t *testing.T) {
 	assert.Contains(t, mermaid, "   task2 --> task1\n")
 	assert.Contains(t, mermaid, "   task3 --> task1\n")
 	assert.Contains(t, mermaid, "   task3 --> task2\n")
-	assert.Contains(t, mermaid, "   task1 --> id([fa:fa-circle])\n")
-	assert.Contains(t, mermaid, "   task4 --> id([fa:fa-circle])\n")
+	assert.Contains(t, mermaid, "   task1 --> stop([fa:fa-circle])\n")
+	assert.Contains(t, mermaid, "   task4 --> stop([fa:fa-circle])\n")
+	assert.Contains(t, mermaid, "   start([fa:fa-circle]) --> task3\n")
+	assert.Contains(t, mermaid, "   start([fa:fa-circle]) --> task4\n")
 }
 
 func TestTaskGraphToMermaidWithTaskRef(t *testing.T) {
@@ -179,8 +181,10 @@ func TestTaskGraphToMermaidWithTaskRef(t *testing.T) {
 	assert.Contains(t, mermaid, "   task2(\"task2\n   (taskRef2)\") --> task1(\"task1\n   (taskRef1)\")\n")
 	assert.Contains(t, mermaid, "   task3(\"task3\n   (taskRef3)\") --> task1(\"task1\n   (taskRef1)\")\n")
 	assert.Contains(t, mermaid, "   task3(\"task3\n   (taskRef3)\") --> task2(\"task2\n   (taskRef2)\")\n")
-	assert.Contains(t, mermaid, "   task1(\"task1\n   (taskRef1)\") --> id([fa:fa-circle])\n")
-	assert.Contains(t, mermaid, "   task4(\"task4\n   (taskRef4)\") --> id([fa:fa-circle])\n")
+	assert.Contains(t, mermaid, "   task1(\"task1\n   (taskRef1)\") --> stop([fa:fa-circle])\n")
+	assert.Contains(t, mermaid, "   task4(\"task4\n   (taskRef4)\") --> stop([fa:fa-circle])\n")
+	assert.Contains(t, mermaid, "   start([fa:fa-circle]) --> task3(\"task3\n   (taskRef3)\")\n")
+	assert.Contains(t, mermaid, "   start([fa:fa-circle]) --> task4(\"task4\n   (taskRef4)\")\n")
 }
 
 func TestFormatFunc(t *testing.T) {