Parse contexts that have Suite interface in func definition

builder.go

@@ -6,10 +6,10 @@ import (
 	"go/format"
 	"go/parser"
 	"go/token"
-	"log"
 	"os"
 	"path/filepath"
 	"strings"
+	"text/template"
 
 	"golang.org/x/tools/imports"
 )
@@ -21,13 +21,20 @@ import (
 )
 
 func main() {
-	godog.Run()
+	suite := &GodogSuite{
+		steps: make(map[*regexp.Regexp]StepHandler),
+	}
+	{{range $c := .Contexts}}
+		{{$c}}(suite)
+	{{end}}
+	suite.Run()
 }
 `
 
 type builder struct {
-	files map[string]*ast.File
+	files    map[string]*ast.File
-	fset  *token.FileSet
+	fset     *token.FileSet
+	Contexts []string
 }
 
 func newBuilder() *builder {
@@ -43,8 +50,9 @@ func (b *builder) parseFile(path string) error {
 		return err
 	}
 	b.deleteMainFunc(f)
+	b.registerSteps(f)
 	b.deleteImports(f)
-	b.files[f.Name.Name] = f
+	b.files[path] = f
 	return nil
 }
 
@@ -78,10 +86,33 @@ func (b *builder) deleteMainFunc(f *ast.File) {
 	f.Decls = decls
 }
 
+func (b *builder) registerSteps(f *ast.File) {
+	for _, d := range f.Decls {
+		fun, ok := d.(*ast.FuncDecl)
+		if !ok {
+			continue
+		}
+		for _, param := range fun.Type.Params.List {
+			ident, ok := param.Type.(*ast.Ident)
+			if !ok {
+				continue
+			}
+			if ident.Name == "godog.Suite" || f.Name.Name == "godog" && ident.Name == "Suite" {
+				b.Contexts = append(b.Contexts, fun.Name.Name)
+			}
+		}
+	}
+}
+
 func (b *builder) merge() (*ast.File, error) {
-	f, err := parser.ParseFile(b.fset, "", mainTpl, 0)
+	var buf bytes.Buffer
+	t := template.Must(template.New("main").Parse(mainTpl))
+	if err := t.Execute(&buf, b); err != nil {
+		return nil, err
+	}
+
+	f, err := parser.ParseFile(b.fset, "", &buf, 0)
 	if err != nil {
-		log.Println("fail here")
 		return nil, err
 	}
 	b.deleteImports(f)

runner.go

@@ -1,14 +0,0 @@
-package godog
-
-import (
-	"flag"
-	"log"
-)
-
-func Run() {
-	if !flag.Parsed() {
-		flag.Parse()
-	}
-	log.Println("running godoc, num registered steps:", len(stepHandlers), "color test:", cl("red", red))
-	log.Println("will read features in path:", cl(cfg.featuresPath, yellow))
-}

steps_test.go

@@ -5,10 +5,18 @@ import (
 	"regexp"
 )
 
-func init() {
+//func init() {
+//f := StepHandlerFunc(func(args ...interface{}) error {
+//log.Println("step triggered")
+//return nil
+//})
+//Step(regexp.MustCompile("hello"), f)
+//}
+
+func SomeContext(g Suite) {
 	f := StepHandlerFunc(func(args ...interface{}) error {
 		log.Println("step triggered")
 		return nil
 	})
-	Step(regexp.MustCompile("hello"), f)
+	g.Step(regexp.MustCompile("hello"), f)
 }

suite.go

@@ -0,0 +1,22 @@
+package godog
+
+import (
+	"log"
+	"regexp"
+)
+
+type Suite interface {
+	Step(exp *regexp.Regexp, h StepHandler)
+}
+
+type GodogSuite struct {
+	steps map[*regexp.Regexp]StepHandler
+}
+
+func (s *GodogSuite) Step(exp *regexp.Regexp, h StepHandler) {
+	s.steps[exp] = h
+}
+
+func (s *GodogSuite) Run() {
+	log.Println("running godoc, num registered steps:", len(s.steps))
+}