update test code

Signed-off-by: Andres Taylor <[email protected]>

go/tools/asthelpergen/ast_path_gen.go

@@ -49,7 +49,16 @@ func newPathGen(pkgname string) *pathGen {
 }
 
 func (p *pathGen) genFile() (string, *jen.File) {
-	p.close()
+	p.file.ImportName("fmt", "fmt")
+
+	// Declare the ASTStep type with underlying type uint16
+	p.file.Add(jen.Type().Id("ASTStep").Uint16())
+
+	// Add the const block
+	p.file.Add(p.buildConstWithEnum())
+
+	// Add the ASTStep#DebugString() method to the file
+	p.file.Add(p.debugString())
 	return "ast_path.go", p.file
 }
 
@@ -109,19 +118,6 @@ func (p *pathGen) basicMethod(t types.Type, basic *types.Basic, spi generatorSPI
 	return nil
 }
 
-func (p *pathGen) close() {
-	p.file.ImportName("fmt", "fmt")
-
-	// Declare the ASTStep type with underlying type uint16
-	p.file.Add(jen.Type().Id("ASTStep").Uint16())
-
-	// Add the const block
-	p.file.Add(p.buildConstWithEnum())
-
-	// Add the DebugString() method to the file
-	p.file.Add(p.debugString())
-}
-
 func (p *pathGen) debugString() *jen.Statement {
 	var switchCases []jen.Code
 

go/tools/asthelpergen/asthelpergen_test.go

@@ -21,6 +21,8 @@ import (
 	"strings"
 	"testing"
 
+	"vitess.io/vitess/go/tools/codegen"
+
 	"github.com/stretchr/testify/require"
 )
 
@@ -45,3 +47,20 @@ func TestFullGeneration(t *testing.T) {
 		require.False(t, applyIdx == 0 && cloneIdx == 0, "file doesn't contain expected contents")
 	}
 }
+
+func TestRecreateAllFiles(t *testing.T) {
+	t.Skip("This test recreates all files in the integration directory. It should only be run when the ASTHelperGen code has changed.")
+	result, err := GenerateASTHelpers(&Options{
+		Packages:      []string{"./integration/..."},
+		RootInterface: "vitess.io/vitess/go/tools/asthelpergen/integration.AST",
+		Clone: CloneOptions{
+			Exclude: []string{"*NoCloneType"},
+		},
+	})
+	require.NoError(t, err)
+
+	for fullPath, file := range result {
+		err := codegen.SaveJenFile(fullPath, file)
+		require.NoError(t, err)
+	}
+}