♻️ Refactor recurse function into visitor

clevyr · May 6, 2022 · b463169 · b463169
1 parent 7c00a3a
commit b463169
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -148,7 +148,7 @@ func templateReader(r io.Reader) ([]byte, error) {
 			buf.Write([]byte("---\n"))
 		}
 
-		if err := template.RecurseNode(conf, &node); err != nil {
+		if err := template.VisitNodes(conf, template.LineComment, &node); err != nil {
 			return buf.Bytes(), err
 		}
 

diff --git a/pkg/template/line_comment.go b/pkg/template/line_comment.go
@@ -15,14 +15,16 @@ func init() {
 	funcMap["tag"] = DockerTag
 }
 
-func RecurseNode(conf config.Config, node *yaml.Node) error {
+type Visitor func(conf config.Config, node *yaml.Node) error
+
+func VisitNodes(conf config.Config, visit Visitor, node *yaml.Node) error {
 	if len(node.Content) == 0 {
 		if err := LineComment(conf, node); err != nil {
 			return err
 		}
 	} else {
 		for _, node := range node.Content {
-			if err := RecurseNode(conf, node); err != nil {
+			if err := VisitNodes(conf, visit, node); err != nil {
 				return err
 			}
 		}

diff --git a/pkg/template/line_comment_test.go b/pkg/template/line_comment_test.go
@@ -37,17 +37,17 @@ func TestRecurseNode(t *testing.T) {
 			var node yaml.Node
 			_ = yaml.Unmarshal([]byte(tt.args.input), &node)
 
-			if err := RecurseNode(tt.args.conf, &node); err != nil {
+			if err := VisitNodes(tt.args.conf, LineComment, &node); err != nil {
 				if (err != nil) != tt.wantErr {
-					t.Errorf("RecurseNode() error = %v, wantErr %v", err, tt.wantErr)
+					t.Errorf("VisitNodes() error = %v, wantErr %v", err, tt.wantErr)
 				}
 				return
 			}
 
 			got, _ := yaml.Marshal(&node)
 			got = bytes.TrimRight(got, "\n")
 			if string(got) != tt.want {
-				t.Errorf("RecurseNode() = %v, want %v", string(got), tt.want)
+				t.Errorf("VisitNodes() = %v, want %v", string(got), tt.want)
 			}
 		})
 	}