From 3f8ae6e84a7d9eca99d3bcadefba7b40761f7e0d Mon Sep 17 00:00:00 2001
From: Samsondeen Dare <samsondeen.dare@hashicorp.com>
Date: Wed, 8 Jan 2025 20:06:13 +0100
Subject: [PATCH] fix tests

---
 internal/dag/dag.go       |  5 +++--
 internal/dag/walk.go      | 14 ++++++++++++++
 internal/dag/walk_test.go | 21 ++++++++++-----------
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/internal/dag/dag.go b/internal/dag/dag.go
index de5476c5cee4..08c97b8a3e09 100644
--- a/internal/dag/dag.go
+++ b/internal/dag/dag.go
@@ -278,8 +278,9 @@ func (g *AcyclicGraph) Cycles() [][]Vertex {
 // This will walk nodes in parallel if it can. The resulting diagnostics
 // contains problems from all graphs visited, in no particular order.
 func (g *AcyclicGraph) Walk(cb WalkFunc) tfdiags.Diagnostics {
-	ctx, cancel := context.WithCancelCause(context.Background())
-	w := &Walker{Callback: cb, Reverse: true, walkContext: ctx, walkContextCancel: cancel}
+	w := NewWalker(cb, func(w *Walker) {
+		w.Reverse = true
+	})
 	g.walker = w
 
 	w.Update(g)
diff --git a/internal/dag/walk.go b/internal/dag/walk.go
index 3ce57307f2f5..b46b9cfc4fb5 100644
--- a/internal/dag/walk.go
+++ b/internal/dag/walk.go
@@ -80,6 +80,20 @@ func (w *Walker) init() {
 	}
 }
 
+// NewWalker creates a new walker with the given callback function.
+func NewWalker(cb WalkFunc, opts ...func(*Walker)) *Walker {
+	ctx, cancel := context.WithCancelCause(context.Background())
+	w := &Walker{
+		Callback:          cb,
+		walkContext:       ctx,
+		walkContextCancel: cancel,
+	}
+	for _, opt := range opts {
+		opt(w)
+	}
+	return w
+}
+
 type walkerVertex struct {
 	// These should only be set once on initialization and never written again.
 	// They are not protected by a lock since they don't need to be since
diff --git a/internal/dag/walk_test.go b/internal/dag/walk_test.go
index ddb3b8055b1e..ef2353fd96c8 100644
--- a/internal/dag/walk_test.go
+++ b/internal/dag/walk_test.go
@@ -22,7 +22,7 @@ func TestWalker_basic(t *testing.T) {
 	// Run it a bunch of times since it is timing dependent
 	for i := 0; i < 50; i++ {
 		var order []interface{}
-		w := &Walker{Callback: walkCbRecord(&order)}
+		w := NewWalker(walkCbRecord(&order))
 		w.Update(&g)
 
 		// Wait
@@ -47,7 +47,7 @@ func TestWalker_updateNilGraph(t *testing.T) {
 	// Run it a bunch of times since it is timing dependent
 	for i := 0; i < 50; i++ {
 		var order []interface{}
-		w := &Walker{Callback: walkCbRecord(&order)}
+		w := NewWalker(walkCbRecord(&order))
 		w.Update(&g)
 		w.Update(nil)
 
@@ -83,7 +83,7 @@ func TestWalker_error(t *testing.T) {
 		return recordF(v)
 	}
 
-	w := &Walker{Callback: cb}
+	w := NewWalker(cb)
 	w.Update(&g)
 
 	// Wait
@@ -110,7 +110,6 @@ func TestWalker_newVertex(t *testing.T) {
 	done2 := make(chan int)
 
 	// Build a callback that notifies us when 2 has been walked
-	var w *Walker
 	cb := func(v Vertex) tfdiags.Diagnostics {
 		if v == 2 {
 			defer close(done2)
@@ -119,7 +118,7 @@ func TestWalker_newVertex(t *testing.T) {
 	}
 
 	// Add the initial vertices
-	w = &Walker{Callback: cb}
+	w := NewWalker(cb)
 	w.Update(&g)
 
 	// if 2 has been visited, the walk is complete so far
@@ -155,7 +154,7 @@ func TestWalker_removeVertex(t *testing.T) {
 	var order []interface{}
 	recordF := walkCbRecord(&order)
 
-	var w *Walker
+	w := NewWalker(nil)
 	cb := func(v Vertex) tfdiags.Diagnostics {
 		if v == 1 {
 			g.Remove(2)
@@ -166,7 +165,7 @@ func TestWalker_removeVertex(t *testing.T) {
 	}
 
 	// Add the initial vertices
-	w = &Walker{Callback: cb}
+	w.Callback = cb
 	w.Update(&g)
 
 	// Wait
@@ -191,7 +190,7 @@ func TestWalker_newEdge(t *testing.T) {
 	var order []interface{}
 	recordF := walkCbRecord(&order)
 
-	var w *Walker
+	w := NewWalker(nil)
 	cb := func(v Vertex) tfdiags.Diagnostics {
 		// record where we are first, otherwise the Updated vertex may get
 		// walked before the first visit.
@@ -206,7 +205,7 @@ func TestWalker_newEdge(t *testing.T) {
 	}
 
 	// Add the initial vertices
-	w = &Walker{Callback: cb}
+	w.Callback = cb
 	w.Update(&g)
 
 	// Wait
@@ -241,7 +240,7 @@ func TestWalker_removeEdge(t *testing.T) {
 	// forcing 2 before 3 via the callback (and not the graph). If
 	// 2 cannot execute before 3 (edge removal is non-functional), then
 	// this test will timeout.
-	var w *Walker
+	w := NewWalker(nil)
 	gateCh := make(chan struct{})
 	cb := func(v Vertex) tfdiags.Diagnostics {
 		t.Logf("visit vertex %#v", v)
@@ -274,7 +273,7 @@ func TestWalker_removeEdge(t *testing.T) {
 	}
 
 	// Add the initial vertices
-	w = &Walker{Callback: cb}
+	w.Callback = cb
 	w.Update(&g)
 
 	// Wait