Refactor tests in Flux

refactor_tests.go

package main

import (
	"fmt"
	"github.com/influxdata/flux/ast"
	"github.com/influxdata/flux/parser"
	"io/ioutil"
	"os"
	"path/filepath"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "provide testdata folder, please")
		os.Exit(1)
	}

	testDataDir := os.Args[1]
	baseDir, _ := filepath.Split(testDataDir)
	newTestDataDir := filepath.Join(baseDir, "testdata_new")

	fluxFiles, err := filepath.Glob(testDataDir + "/*.flux")
	if err != nil {
		panic(err)
	}

	err = os.RemoveAll(newTestDataDir)
	if err != nil {
		panic(err)
	}
	err = os.Mkdir(newTestDataDir, 0755)
	if err != nil {
		panic(err)
	}

	for _, fluxFile := range fluxFiles {
		ext := filepath.Ext(fluxFile)
		prefix := fluxFile[0 : len(fluxFile)-len(ext)]
		_, caseName := filepath.Split(prefix)
		fName := filepath.Join(newTestDataDir, caseName+ext)

		fmt.Printf("\nEditing %s...\n", fName)

		// read file
		flux, err := ioutil.ReadFile(fluxFile)
		if err != nil {
			panic(err)
		}

		// get AST
		astPkg := parser.ParseSource(string(flux))
		if ast.Check(astPkg) > 0 {
			panic(ast.GetError(astPkg))
		}

		for i, node := range astPkg.Files[0].Body {
			if expr, ok := node.(*ast.ExpressionStatement); ok {
				if fn, ok := expr.Expression.(*ast.CallExpression); ok {
					if callee, ok := fn.Callee.(*ast.MemberExpression); ok {
						if pack, ok := callee.Object.(*ast.Identifier); ok {
							if fnName, ok := callee.Property.(*ast.Identifier); ok {
								if pack.Name == "testing" && fnName.Name == "run" {
									// ok, change stuff
									arg := fn.Arguments[0].(*ast.ObjectExpression)
									var name string
									input := &ast.Property{
										Key: &ast.Identifier{Name: "input"},
									}
									want := &ast.Property{
										Key: &ast.Identifier{Name: "want"},
									}
									fn := &ast.Property{
										Key: &ast.Identifier{Name: "fn"},
									}
									for _, prop := range arg.Properties {
										switch prop.Key.Key() {
										case "name":
											name = prop.Value.(*ast.StringLiteral).Value
										case "input":
											input.Value = prop.Value
										case "want":
											want.Value = prop.Value
										case "testFn":
											fn.Value = prop.Value
										}
									}

									testStmt := &ast.TestStatement{
										Assignment: &ast.VariableAssignment{
											ID: &ast.Identifier{Name: name},
											Init: &ast.ObjectExpression{
												Properties: []*ast.Property{input, want, fn},
											},
										},
									}

									fmt.Println(ast.Format(testStmt))

									astPkg.Files[0].Body[i] = testStmt
								}
							}
						}
					}
				}
			}
		}

		fmtAst := ast.Format(astPkg)

		err = ioutil.WriteFile(fName, []byte(fmtAst), 0644)
		if err != nil {
			fmt.Fprint(os.Stderr, err)
		} else {
			fmt.Printf("%s: done\n", fName)
		}
	}
}