From 3546eda9a2df11aed10ee7a6b760ba515ae989f0 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Thu, 27 Feb 2025 14:24:47 +0100
Subject: [PATCH] test: add test for Visitable Rewrite and Walk

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 .../integration/integration_rewriter_test.go  | 19 +++++++++++
 .../integration/integration_visit_test.go     | 33 +++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/go/tools/asthelpergen/integration/integration_rewriter_test.go b/go/tools/asthelpergen/integration/integration_rewriter_test.go
index 067417ecbf9..13eb8916001 100644
--- a/go/tools/asthelpergen/integration/integration_rewriter_test.go
+++ b/go/tools/asthelpergen/integration/integration_rewriter_test.go
@@ -368,6 +368,25 @@ func TestRefSliceContainerReplace(t *testing.T) {
 	}, ast)
 }
 
+func TestVisitableRewrite(t *testing.T) {
+	leaf := &Leaf{v: 1}
+	visitable := &testVisitable{inner: leaf}
+	refContainer := &RefContainer{ASTType: visitable}
+
+	tv := &rewriteTestVisitor{}
+
+	_ = Rewrite(refContainer, tv.pre, tv.post)
+
+	tv.assertEquals(t, []step{
+		Pre{refContainer},
+		Pre{visitable},
+		Pre{leaf},
+		Post{leaf},
+		Post{visitable},
+		Post{refContainer},
+	})
+}
+
 type step interface {
 	String() string
 }
diff --git a/go/tools/asthelpergen/integration/integration_visit_test.go b/go/tools/asthelpergen/integration/integration_visit_test.go
index 05773e73b8c..e7b16fa7860 100644
--- a/go/tools/asthelpergen/integration/integration_visit_test.go
+++ b/go/tools/asthelpergen/integration/integration_visit_test.go
@@ -151,6 +151,39 @@ func TestVisitInterfaceSlice(t *testing.T) {
 	})
 }
 
+// testVisitable is declared in a test package, so it lacks the normal generated AST helper methods.
+// Instead, we lean on the Visitable interface to make sure we can visit nodes inside it.
+type testVisitable struct {
+	inner AST
+}
+
+func (t *testVisitable) String() string {
+	return t.inner.String()
+}
+
+func (t *testVisitable) VisitThis() AST {
+	return t.inner
+}
+
+var _ Visitable = (*testVisitable)(nil)
+
+func TestVisitableVisit(t *testing.T) {
+	leaf := &Leaf{v: 1}
+	visitable := &testVisitable{inner: leaf}
+	refContainer := &RefContainer{ASTType: visitable}
+
+	tv := &testVisitor{}
+
+	require.NoError(t,
+		VisitAST(refContainer, tv.visit))
+
+	tv.assertVisitOrder(t, []AST{
+		refContainer,
+		visitable,
+		leaf,
+	})
+}
+
 func (tv *testVisitor) assertVisitOrder(t *testing.T, expected []AST) {
 	t.Helper()
 	var lines []string