test: fix test expectations now that we use LeafSlice

Signed-off-by: Manan Gupta <manan@planetscale.com>
planetscale · Feb 11, 2025 · 59c422f · 59c422f
1 parent 061e8ed
commit 59c422f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/go/tools/asthelpergen/integration/integration_rewriter_test.go b/go/tools/asthelpergen/integration/integration_rewriter_test.go
@@ -105,7 +105,8 @@ func TestRewriteVisitValueSliceContainer(t *testing.T) {
 	leaf2 := &Leaf{2}
 	leaf3 := &Leaf{3}
 	leaf4 := &Leaf{4}
-	container := ValueSliceContainer{ASTElements: []AST{leaf1, leaf2}, ASTImplementationElements: []*Leaf{leaf3, leaf4}}
+	ls := LeafSlice{leaf3, leaf4}
+	container := ValueSliceContainer{ASTElements: []AST{leaf1, leaf2}, ASTImplementationElements: ls}
 	containerContainer := ValueSliceContainer{ASTElements: []AST{container}}
 
 	tv := &rewriteTestVisitor{}
@@ -119,10 +120,12 @@ func TestRewriteVisitValueSliceContainer(t *testing.T) {
 		Post{leaf1},
 		Pre{leaf2},
 		Post{leaf2},
+		Pre{ls},
 		Pre{leaf3},
 		Post{leaf3},
 		Pre{leaf4},
 		Post{leaf4},
+		Post{ls},
 		Post{container},
 		Post{containerContainer},
 	})
@@ -402,6 +405,7 @@ func (tv *rewriteTestVisitor) assertEquals(t *testing.T, expected []step) {
 	expectedSize := len(expected)
 	for i, step := range tv.walk {
 		t.Run(fmt.Sprintf("step %d", i), func(t *testing.T) {
+			t.Helper()
 			if expectedSize <= i {
 				t.Fatalf("❌️ - Expected less elements %v", tv.walk[i:])
 			} else {

diff --git a/go/tools/asthelpergen/integration/integration_visit_test.go b/go/tools/asthelpergen/integration/integration_visit_test.go
@@ -100,7 +100,8 @@ func TestVisitValueSliceContainer(t *testing.T) {
 	leaf2 := &Leaf{2}
 	leaf3 := &Leaf{3}
 	leaf4 := &Leaf{4}
-	container := ValueSliceContainer{ASTElements: []AST{leaf1, leaf2}, ASTImplementationElements: []*Leaf{leaf3, leaf4}}
+	ls := LeafSlice{leaf3, leaf4}
+	container := ValueSliceContainer{ASTElements: []AST{leaf1, leaf2}, ASTImplementationElements: ls}
 	containerContainer := ValueSliceContainer{ASTElements: []AST{container}}
 
 	tv := &testVisitor{}
@@ -113,6 +114,7 @@ func TestVisitValueSliceContainer(t *testing.T) {
 		container,
 		leaf1,
 		leaf2,
+		ls,
 		leaf3,
 		leaf4,
 	})