fix context import when package is not godog

builder.go

@@ -85,19 +85,24 @@ func (b *builder) deleteMainFunc(f *ast.File) {
 
 func (b *builder) registerSteps(f *ast.File) {
 	for _, d := range f.Decls {
-		fun, ok := d.(*ast.FuncDecl)
-		if !ok {
-			continue
-		}
+		switch fun := d.(type) {
+		case *ast.FuncDecl:
 			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" {
+				switch expr := param.Type.(type) {
+				case *ast.SelectorExpr:
+					switch x := expr.X.(type) {
+					case *ast.Ident:
+						if x.Name == "godog" && expr.Sel.Name == "Suite" {
 							b.Contexts = append(b.Contexts, fun.Name.Name)
 						}
 					}
+				case *ast.Ident:
+					if expr.Name == "Suite" {
+						b.Contexts = append(b.Contexts, fun.Name.Name)
+					}
+				}
+			}
+		}
 	}
 }
 

config.go

@@ -1,7 +1,6 @@
 package godog
 
 import (
-	"flag"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -17,12 +16,6 @@ type config struct {
 	formatterName string
 }
 
-func init() {
-	// @TODO: colorize flag help output
-	flag.StringVar(&cfg.featuresPath, "features", "features", "Path to feature files")
-	flag.StringVar(&cfg.formatterName, "formatter", "pretty", "Formatter name")
-}
-
 func (c config) validate() error {
 	inf, err := os.Stat(c.featuresPath)
 	if err != nil {

suite.go

@@ -49,6 +49,9 @@ type suite struct {
 // interface. The instance is passed around to all
 // context initialization functions from *_test.go files
 func New() *suite {
+	// @TODO: colorize flag help output
+	flag.StringVar(&cfg.featuresPath, "features", "features", "Path to feature files")
+	flag.StringVar(&cfg.formatterName, "formatter", "pretty", "Formatter name")
 	if !flag.Parsed() {
 		flag.Parse()
 	}