diff --git a/builder.go b/builder.go
index bc008fc..f0db029 100644
--- a/builder.go
+++ b/builder.go
@@ -85,17 +85,22 @@ 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
-		}
-		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)
+		switch fun := d.(type) {
+		case *ast.FuncDecl:
+			for _, param := range fun.Type.Params.List {
+				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)
+					}
+				}
 			}
 		}
 	}
diff --git a/config.go b/config.go
index 1359004..0dca616 100644
--- a/config.go
+++ b/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 {
diff --git a/suite.go b/suite.go
index 4017d57..9d57a28 100644
--- a/suite.go
+++ b/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()
 	}