reorganize test package builder

ast.go

@@ -0,0 +1,109 @@
+package godog
+
+import (
+	"go/ast"
+	"go/build"
+	"go/token"
+	"path"
+	"strings"
+)
+
+func removeUnusedImports(f *ast.File) {
+	used := usedPackages(f)
+	isUsed := func(p string) bool {
+		for _, ref := range used {
+			if p == ref {
+				return true
+			}
+		}
+		return p == "_"
+	}
+	var decls []ast.Decl
+	for _, d := range f.Decls {
+		gen, ok := d.(*ast.GenDecl)
+		if ok && gen.Tok == token.IMPORT {
+			var specs []ast.Spec
+			for _, spec := range gen.Specs {
+				impspec := spec.(*ast.ImportSpec)
+				ipath := strings.Trim(impspec.Path.Value, `\"`)
+				check := importPathToName(ipath)
+				if impspec.Name != nil {
+					check = impspec.Name.Name
+				}
+
+				if isUsed(check) {
+					specs = append(specs, spec)
+				}
+			}
+
+			if len(specs) == 0 {
+				continue
+			}
+			gen.Specs = specs
+		}
+		decls = append(decls, d)
+	}
+	f.Decls = decls
+}
+
+func deleteTestMainFunc(f *ast.File) {
+	var decls []ast.Decl
+	var hadMain bool
+	for _, d := range f.Decls {
+		fun, ok := d.(*ast.FuncDecl)
+		if !ok {
+			decls = append(decls, d)
+			continue
+		}
+		if fun.Name.Name != "TestMain" {
+			decls = append(decls, fun)
+		} else {
+			hadMain = true
+		}
+	}
+	f.Decls = decls
+
+	if hadMain {
+		removeUnusedImports(f)
+	}
+}
+
+type visitFn func(node ast.Node) ast.Visitor
+
+func (fn visitFn) Visit(node ast.Node) ast.Visitor {
+	return fn(node)
+}
+
+func usedPackages(f *ast.File) []string {
+	var refs []string
+	var visitor visitFn
+	visitor = visitFn(func(node ast.Node) ast.Visitor {
+		if node == nil {
+			return visitor
+		}
+		switch v := node.(type) {
+		case *ast.SelectorExpr:
+			xident, ok := v.X.(*ast.Ident)
+			if !ok {
+				break
+			}
+			if xident.Obj != nil {
+				// if the parser can resolve it, it's not a package ref
+				break
+			}
+			refs = append(refs, xident.Name)
+		}
+		return visitor
+	})
+	ast.Walk(visitor, f)
+	return refs
+}
+
+// importPathToName finds out the actual package name, as declared in its .go files.
+// If there's a problem, it falls back to using importPathToNameBasic.
+func importPathToName(importPath string) (packageName string) {
+	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
+		return buildPkg.Name
+	}
+	return path.Base(importPath)
+}

ast_test.go

@@ -0,0 +1,117 @@
+package godog
+
+var builderMainFile = `
+package main
+import "fmt"
+func main() {
+	fmt.Println("hello")
+}`
+
+var builderTestMainFile = `
+package main
+import (
+	"fmt"
+	"testing"
+	"os"
+)
+func TestMain(m *testing.M) {
+	fmt.Println("hello")
+	os.Exit(0)
+}`
+
+var builderPackAliases = `
+package main
+import (
+	"testing"
+	a "fmt"
+	b "fmt"
+)
+func TestMain(m *testing.M) {
+	a.Println("a")
+	b.Println("b")
+}`
+
+var builderAnonymousImport = `
+package main
+import (
+	"testing"
+	_ "github.com/go-sql-driver/mysql"
+)
+func TestMain(m *testing.M) {
+}`
+
+var builderContextSrc = `
+package main
+import (
+	"github.com/DATA-DOG/godog"
+)
+func myContext(s *godog.Suite) {
+}
+`
+
+var builderLibrarySrc = `
+package lib
+import "fmt"
+func test() {
+	fmt.Println("hello")
+}
+`
+
+var builderInternalPackageSrc = `
+package godog
+import "fmt"
+func test() {
+	fmt.Println("hello")
+}
+`
+
+// func builderProcess(src string, t *testing.T) string {
+// 	fset := token.NewFileSet()
+// 	f, err := parser.ParseFile(fset, "", []byte(builderTestMainFile), 0)
+// 	if err != nil {
+// 		t.Fatalf("unexpected error while parsing ast: %v", err)
+// 	}
+
+// 	deleteTestMainFunc(f)
+
+// 	var buf strings.Buffer
+// 	if err := format.Node(&buf, fset, node); err != nil {
+// 		return err
+// 	}
+// }
+
+// func TestShouldCleanTestMainFromSimpleTestFile(t *testing.T) {
+
+// 	b := newBuilderSkel()
+// 	err := b.registerMulti([]string{
+// 		builderMainFile, builderPackAliases, builderAnonymousImport,
+// 	})
+// 	if err != nil {
+// 		t.Fatalf("unexpected error: %s", err)
+// 	}
+
+// 	data, err := b.merge()
+// 	if err != nil {
+// 		t.Fatalf("unexpected error: %s", err)
+// 	}
+// 	expected := `package main
+// import (
+// 	a "fmt"
+// 	b "fmt"
+// 	"github.com/DATA-DOG/godog"
+// 	_ "github.com/go-sql-driver/mysql"
+// )
+// func main() {
+// 	godog.Run(func(suite *godog.Suite) {
+// 	})
+// }
+// func Tester() {
+// 	a.Println("a")
+// 	b.Println("b")
+// }`
+
+// 	actual := string(data)
+// 	if b.cleanSpacing(expected) != b.cleanSpacing(actual) {
+// 		t.Fatalf("expected output does not match: %s", actual)
+// 	}
+// }

builder.go

@@ -3,49 +3,16 @@ package godog
 import (
 	"bytes"
 	"go/ast"
-	"go/build"
 	"go/format"
 	"go/parser"
 	"go/token"
 	"os"
-	"path"
 	"path/filepath"
 	"strings"
 	"text/template"
 )
 
-type builder struct {
+var runnerTemplate = template.Must(template.New("main").Parse(`package main
-	files     map[string]*ast.File
-	fset      *token.FileSet
-	Contexts  []string
-	Internal  bool
-	SuiteName string
-	tpl       *template.Template
-
-	imports []*ast.ImportSpec
-}
-
-func (b *builder) hasImport(a *ast.ImportSpec) bool {
-	for _, b := range b.imports {
-		var aname, bname string
-		if a.Name != nil {
-			aname = a.Name.Name
-		}
-		if b.Name != nil {
-			bname = b.Name.Name
-		}
-		if bname == aname && a.Path.Value == b.Path.Value {
-			return true
-		}
-	}
-	return false
-}
-
-func newBuilderSkel() *builder {
-	return &builder{
-		files: make(map[string]*ast.File),
-		fset:  token.NewFileSet(),
-		tpl: template.Must(template.New("main").Parse(`package main
 import (
 {{ if not .Internal }}	"github.com/DATA-DOG/godog"{{ end }}
 	"os"
@@ -61,54 +28,13 @@ func TestMain(m *testing.M) {
 		{{end}}
 	})
 	os.Exit(status)
-}`)),
+}`))
-	}
-}
 
-func doBuild(buildPath, dir string) error {
+type builder struct {
-	b := newBuilderSkel()
+	files     map[string]*ast.File
-	err := filepath.Walk(buildPath, func(path string, file os.FileInfo, err error) error {
+	Contexts  []string
-		if file.IsDir() && file.Name() != "." {
+	Internal  bool
-			return filepath.SkipDir
+	SuiteName string
-		}
-		if err == nil && strings.HasSuffix(path, ".go") {
-			f, err := parser.ParseFile(b.fset, path, nil, 0)
-			if err != nil {
-				return err
-			}
-			b.register(f, file.Name())
-		}
-		return err
-	})
-
-	if err != nil {
-		return err
-	}
-
-	var buf bytes.Buffer
-	if err := b.tpl.Execute(&buf, b); err != nil {
-		return err
-	}
-
-	f, err := parser.ParseFile(b.fset, "", &buf, 0)
-	if err != nil {
-		return err
-	}
-
-	b.files["godog_test.go"] = f
-
-	os.Mkdir(dir, 0755)
-	for name, node := range b.files {
-		f, err := os.Create(filepath.Join(dir, name))
-		if err != nil {
-			return err
-		}
-		if err := format.Node(f, b.fset, node); err != nil {
-			return err
-		}
-	}
-
-	return nil
 }
 
 func (b *builder) register(f *ast.File, name string) {
@@ -117,90 +43,12 @@ func (b *builder) register(f *ast.File, name string) {
 		b.Internal = true
 	}
 	b.SuiteName = f.Name.Name
-	b.deleteMainFunc(f)
+	deleteTestMainFunc(f)
 	f.Name.Name = "main"
 	b.registerContexts(f)
 	b.files[name] = f
 }
 
-func (b *builder) removeUnusedImports(f *ast.File) {
-	used := b.usedPackages(f)
-	isUsed := func(p string) bool {
-		for _, ref := range used {
-			if p == ref {
-				return true
-			}
-		}
-		return p == "_"
-	}
-	var decls []ast.Decl
-	for _, d := range f.Decls {
-		gen, ok := d.(*ast.GenDecl)
-		if ok && gen.Tok == token.IMPORT {
-			var specs []ast.Spec
-			for _, spec := range gen.Specs {
-				impspec := spec.(*ast.ImportSpec)
-				ipath := strings.Trim(impspec.Path.Value, `\"`)
-				check := importPathToName(ipath)
-				if impspec.Name != nil {
-					check = impspec.Name.Name
-				}
-
-				if isUsed(check) {
-					specs = append(specs, spec)
-				}
-			}
-
-			if len(specs) == 0 {
-				continue
-			}
-			gen.Specs = specs
-		}
-		decls = append(decls, d)
-	}
-	f.Decls = decls
-}
-
-func (b *builder) deleteImports(f *ast.File) {
-	var decls []ast.Decl
-	for _, d := range f.Decls {
-		gen, ok := d.(*ast.GenDecl)
-		if ok && gen.Tok == token.IMPORT {
-			for _, spec := range gen.Specs {
-				impspec := spec.(*ast.ImportSpec)
-				if !b.hasImport(impspec) {
-					b.imports = append(b.imports, impspec)
-				}
-			}
-			continue
-		}
-		decls = append(decls, d)
-	}
-	f.Decls = decls
-}
-
-func (b *builder) deleteMainFunc(f *ast.File) {
-	var decls []ast.Decl
-	var hadMain bool
-	for _, d := range f.Decls {
-		fun, ok := d.(*ast.FuncDecl)
-		if !ok {
-			decls = append(decls, d)
-			continue
-		}
-		if fun.Name.Name != "TestMain" {
-			decls = append(decls, fun)
-		} else {
-			hadMain = true
-		}
-	}
-	f.Decls = decls
-
-	if hadMain {
-		b.removeUnusedImports(f)
-	}
-}
-
 func (b *builder) registerContexts(f *ast.File) {
 	for _, d := range f.Decls {
 		switch fun := d.(type) {
@@ -227,56 +75,61 @@ func (b *builder) registerContexts(f *ast.File) {
 	}
 }
 
-type visitFn func(node ast.Node) ast.Visitor
+// Build scans all go files in current directory,
-
+// copies them to temporary build directory.
-func (fn visitFn) Visit(node ast.Node) ast.Visitor {
+// If there is a TestMain func in any of test.go files
-	return fn(node)
+// it removes it and all necessary unused imports related
-}
+// to this function.
-
-func (b *builder) usedPackages(f *ast.File) []string {
-	var refs []string
-	var visitor visitFn
-	visitor = visitFn(func(node ast.Node) ast.Visitor {
-		if node == nil {
-			return visitor
-		}
-		switch v := node.(type) {
-		case *ast.SelectorExpr:
-			xident, ok := v.X.(*ast.Ident)
-			if !ok {
-				break
-			}
-			if xident.Obj != nil {
-				// if the parser can resolve it, it's not a package ref
-				break
-			}
-			refs = append(refs, xident.Name)
-		}
-		return visitor
-	})
-	ast.Walk(visitor, f)
-	return refs
-}
-
-// Build creates a runnable Godog executable file
-// from current package source and test source files.
 //
-// The package files are merged with the help of go/ast into
+// It also looks for any godog suite contexts and registers
-// a single main package file which has a custom
+// them in order to call them on execution.
-// main function to run test suite features.
 //
-// Currently, to manage imports we use "golang.org/x/tools/imports"
+// The test entry point which uses go1.4 TestMain func
-// package, but that may be replaced in order to have
+// is generated from the template above.
-// no external dependencies
 func Build(dir string) error {
-	return doBuild(".", dir)
+	fset := token.NewFileSet()
-}
+	b := &builder{files: make(map[string]*ast.File)}
 
-// importPathToName finds out the actual package name, as declared in its .go files.
+	err := filepath.Walk(".", func(path string, file os.FileInfo, err error) error {
-// If there's a problem, it falls back to using importPathToNameBasic.
+		if file.IsDir() && file.Name() != "." {
-func importPathToName(importPath string) (packageName string) {
+			return filepath.SkipDir
-	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
-		return buildPkg.Name
 		}
-	return path.Base(importPath)
+		if err == nil && strings.HasSuffix(path, ".go") {
+			f, err := parser.ParseFile(fset, path, nil, 0)
+			if err != nil {
+				return err
+			}
+			b.register(f, file.Name())
+		}
+		return err
+	})
+
+	if err != nil {
+		return err
+	}
+
+	var buf bytes.Buffer
+	if err := runnerTemplate.Execute(&buf, b); err != nil {
+		return err
+	}
+
+	f, err := parser.ParseFile(fset, "", &buf, 0)
+	if err != nil {
+		return err
+	}
+
+	b.files["godog_test.go"] = f
+
+	os.Mkdir(dir, 0755)
+	for name, node := range b.files {
+		f, err := os.Create(filepath.Join(dir, name))
+		if err != nil {
+			return err
+		}
+		if err := format.Node(f, fset, node); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }

builder_test.go

@@ -1 +0,0 @@
-package godog