test builder ast functions

ast.go

@@ -8,6 +8,34 @@ import (
 	"strings"
 )
 
+func contexts(f *ast.File) []string {
+	var contexts []string
+	for _, d := range f.Decls {
+		switch fun := d.(type) {
+		case *ast.FuncDecl:
+			for _, param := range fun.Type.Params.List {
+				switch expr := param.Type.(type) {
+				case *ast.StarExpr:
+					switch x := expr.X.(type) {
+					case *ast.Ident:
+						if x.Name == "Suite" {
+							contexts = append(contexts, fun.Name.Name)
+						}
+					case *ast.SelectorExpr:
+						switch t := x.X.(type) {
+						case *ast.Ident:
+							if t.Name == "godog" && x.Sel.Name == "Suite" {
+								contexts = append(contexts, fun.Name.Name)
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+	return contexts
+}
+
 func removeUnusedImports(f *ast.File) {
 	used := usedPackages(f)
 	isUsed := func(p string) bool {
@@ -48,7 +76,7 @@ func removeUnusedImports(f *ast.File) {
 
 func deleteTestMainFunc(f *ast.File) {
 	var decls []ast.Decl
-	var hadMain bool
+	var hadTestMain bool
 	for _, d := range f.Decls {
 		fun, ok := d.(*ast.FuncDecl)
 		if !ok {
@@ -58,12 +86,12 @@ func deleteTestMainFunc(f *ast.File) {
 		if fun.Name.Name != "TestMain" {
 			decls = append(decls, fun)
 		} else {
-			hadMain = true
+			hadTestMain = true
 		}
 	}
 	f.Decls = decls
 
-	if hadMain {
+	if hadTestMain {
 		removeUnusedImports(f)
 	}
 }

ast_context_test.go

@@ -0,0 +1,76 @@
+package godog
+
+import (
+	"go/parser"
+	"go/token"
+	"testing"
+)
+
+var astContextSrc = `package main
+
+import (
+	"github.com/DATA-DOG/godog"
+)
+
+func myContext(s *godog.Suite) {
+}`
+
+var astTwoContextSrc = `package lib
+
+import (
+	"github.com/DATA-DOG/godog"
+)
+
+func apiContext(s *godog.Suite) {
+}
+
+func dbContext(s *godog.Suite) {
+}`
+
+func astContexts(src string, t *testing.T) []string {
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, "", []byte(src), 0)
+	if err != nil {
+		t.Fatalf("unexpected error while parsing ast: %v", err)
+	}
+
+	return contexts(f)
+}
+
+func TestShouldGetSingleContextFromSource(t *testing.T) {
+	actual := astContexts(astContextSrc, t)
+	expect := []string{"myContext"}
+
+	if len(actual) != len(expect) {
+		t.Fatalf("number of found contexts do not match, expected %d, but got %d", len(expect), len(actual))
+	}
+
+	for i, c := range expect {
+		if c != actual[i] {
+			t.Fatalf("expected context '%s' at pos %d, but got: '%s'", c, i, actual[i])
+		}
+	}
+}
+
+func TestShouldGetTwoContextsFromSource(t *testing.T) {
+	actual := astContexts(astTwoContextSrc, t)
+	expect := []string{"apiContext", "dbContext"}
+
+	if len(actual) != len(expect) {
+		t.Fatalf("number of found contexts do not match, expected %d, but got %d", len(expect), len(actual))
+	}
+
+	for i, c := range expect {
+		if c != actual[i] {
+			t.Fatalf("expected context '%s' at pos %d, but got: '%s'", c, i, actual[i])
+		}
+	}
+}
+
+func TestShouldNotFindAnyContextsInEmptyFile(t *testing.T) {
+	actual := astContexts(`package main`, t)
+
+	if len(actual) != 0 {
+		t.Fatalf("expected no contexts to be found, but there was some: %v", actual)
+	}
+}

ast_test.go

@@ -1,117 +1,162 @@
 package godog
 
-var builderMainFile = `
+import (
-package main
+	"bytes"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"strings"
+	"testing"
+)
+
+var astMainFile = `package main
+
 import "fmt"
+
 func main() {
 	fmt.Println("hello")
 }`
 
-var builderTestMainFile = `
+var astNormalFile = `package main
-package main
+
+import "fmt"
+
+func hello() {
+	fmt.Println("hello")
+}`
+
+var astTestMainFile = `package main
+
 import (
 	"fmt"
 	"testing"
 	"os"
 )
+
 func TestMain(m *testing.M) {
 	fmt.Println("hello")
 	os.Exit(0)
 }`
 
-var builderPackAliases = `
+var astPackAliases = `package main
-package main
+
 import (
 	"testing"
 	a "fmt"
 	b "fmt"
 )
+
 func TestMain(m *testing.M) {
 	a.Println("a")
 	b.Println("b")
 }`
 
-var builderAnonymousImport = `
+var astAnonymousImport = `package main
-package main
+
 import (
 	"testing"
 	_ "github.com/go-sql-driver/mysql"
 )
+
 func TestMain(m *testing.M) {
 }`
 
-var builderContextSrc = `
+var astLibrarySrc = `package lib
-package main
+
+import "fmt"
+
+func test() {
+	fmt.Println("hello")
+}`
+
+var astInternalPackageSrc = `package godog
+
+import "fmt"
+
+func test() {
+	fmt.Println("hello")
+}`
+
+func astProcess(src string, t *testing.T) string {
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, "", []byte(src), 0)
+	if err != nil {
+		t.Fatalf("unexpected error while parsing ast: %v", err)
+	}
+
+	deleteTestMainFunc(f)
+
+	var buf bytes.Buffer
+	if err := format.Node(&buf, fset, f); err != nil {
+		t.Fatalf("failed to build source file: %v", err)
+	}
+
+	return buf.String()
+}
+
+func TestShouldCleanTestMainFromSimpleTestFile(t *testing.T) {
+	actual := strings.TrimSpace(astProcess(astTestMainFile, t))
+	expect := `package main`
+
+	if actual != expect {
+		t.Fatalf("expected output does not match: %s", actual)
+	}
+}
+
+func TestShouldCleanTestMainFromFileWithPackageAliases(t *testing.T) {
+	actual := strings.TrimSpace(astProcess(astPackAliases, t))
+	expect := `package main`
+
+	if actual != expect {
+		t.Fatalf("expected output does not match: %s", actual)
+	}
+}
+
+func TestShouldNotModifyNormalFile(t *testing.T) {
+	actual := strings.TrimSpace(astProcess(astNormalFile, t))
+	expect := astNormalFile
+
+	if actual != expect {
+		t.Fatalf("expected output does not match: %s", actual)
+	}
+}
+
+func TestShouldNotModifyMainFile(t *testing.T) {
+	actual := strings.TrimSpace(astProcess(astMainFile, t))
+	expect := astMainFile
+
+	if actual != expect {
+		t.Fatalf("expected output does not match: %s", actual)
+	}
+}
+
+func TestShouldMaintainAnonymousImport(t *testing.T) {
+	actual := strings.TrimSpace(astProcess(astAnonymousImport, t))
+	expect := `package main
+
 import (
-	"github.com/DATA-DOG/godog"
+	_ "github.com/go-sql-driver/mysql"
-)
+)`
-func myContext(s *godog.Suite) {
+
+	if actual != expect {
+		t.Fatalf("expected output does not match: %s", actual)
+	}
 }
-`
 
-var builderLibrarySrc = `
+func TestShouldNotModifyLibraryPackageSource(t *testing.T) {
-package lib
+	actual := strings.TrimSpace(astProcess(astLibrarySrc, t))
-import "fmt"
+	expect := astLibrarySrc
-func test() {
+
-	fmt.Println("hello")
+	if actual != expect {
+		t.Fatalf("expected output does not match: %s", actual)
+	}
 }
-`
 
-var builderInternalPackageSrc = `
+func TestShouldNotModifyGodogPackageSource(t *testing.T) {
-package godog
+	actual := strings.TrimSpace(astProcess(astInternalPackageSrc, t))
-import "fmt"
+	expect := astInternalPackageSrc
-func test() {
+
-	fmt.Println("hello")
+	if actual != expect {
+		t.Fatalf("expected output does not match: %s", actual)
+	}
 }
-`
-
-// 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

@@ -12,17 +12,17 @@ import (
 	"text/template"
 )
 
-var runnerTemplate = template.Must(template.New("main").Parse(`package main
+var runnerTemplate = template.Must(template.New("main").Parse(`package {{ .PackageName }}
 import (
-{{ if not .Internal }}	"github.com/DATA-DOG/godog"{{ end }}
+{{ if ne .PackageName "godog" }}	"github.com/DATA-DOG/godog"{{ end }}
 	"os"
 	"testing"
 )
 
-const GodogSuiteName = "{{ .SuiteName }}"
+const GodogSuiteName = "{{ .PackageName }}"
 
 func TestMain(m *testing.M) {
-	status := {{ if not .Internal }}godog.{{ end }}Run(func (suite *{{ if not .Internal }}godog.{{ end }}Suite) {
+	status := {{ if ne .PackageName "godog" }}godog.{{ end }}Run(func (suite *{{ if ne .PackageName "godog" }}godog.{{ end }}Suite) {
 		{{range .Contexts}}
 			{{ . }}(suite)
 		{{end}}
@@ -31,50 +31,19 @@ func TestMain(m *testing.M) {
 }`))
 
 type builder struct {
-	files     map[string]*ast.File
+	files       map[string]*ast.File
-	Contexts  []string
+	Contexts    []string
-	Internal  bool
+	PackageName string
-	SuiteName string
 }
 
 func (b *builder) register(f *ast.File, name string) {
-	// mark godog package as internal
+	b.PackageName = f.Name.Name
-	if f.Name.Name == "godog" && !b.Internal {
-		b.Internal = true
-	}
-	b.SuiteName = f.Name.Name
 	deleteTestMainFunc(f)
-	f.Name.Name = "main"
+	// f.Name.Name = "main"
-	b.registerContexts(f)
+	b.Contexts = append(b.Contexts, contexts(f)...)
 	b.files[name] = f
 }
 
-func (b *builder) registerContexts(f *ast.File) {
-	for _, d := range f.Decls {
-		switch fun := d.(type) {
-		case *ast.FuncDecl:
-			for _, param := range fun.Type.Params.List {
-				switch expr := param.Type.(type) {
-				case *ast.StarExpr:
-					switch x := expr.X.(type) {
-					case *ast.Ident:
-						if x.Name == "Suite" {
-							b.Contexts = append(b.Contexts, fun.Name.Name)
-						}
-					case *ast.SelectorExpr:
-						switch t := x.X.(type) {
-						case *ast.Ident:
-							if t.Name == "godog" && x.Sel.Name == "Suite" {
-								b.Contexts = append(b.Contexts, fun.Name.Name)
-							}
-						}
-					}
-				}
-			}
-		}
-	}
-}
-
 // Build scans all go files in current directory,
 // copies them to temporary build directory.
 // If there is a TestMain func in any of test.go files
@@ -94,6 +63,8 @@ func Build(dir string) error {
 		if file.IsDir() && file.Name() != "." {
 			return filepath.SkipDir
 		}
+		// @TODO: maybe should copy all files in root dir (may contain CGO)
+		// or use build.Import go tool, to manage package details
 		if err == nil && strings.HasSuffix(path, ".go") {
 			f, err := parser.ParseFile(fset, path, nil, 0)
 			if err != nil {