[gocomplete] fix test name completion

gocomplete/complete.go

@@ -60,7 +60,7 @@ func main() {
 			"-covermode": complete.PredictSet("set", "count", "atomic"),
 			"-coverpkg":  complete.PredictDirs,
 			"-cpu":       complete.PredictAnything,
-			"-run":       predictTest("test"),
+			"-run":       predictTest("Test", "Example"),
 			"-short":     complete.PredictNothing,
 			"-timeout":   complete.PredictAnything,
 

gocomplete/tests.go

@@ -12,9 +12,14 @@ import (
 	"github.com/posener/complete/match"
 )
 
-func predictTest(testType string) complete.Predicate {
+// predictTest predict test names.
+// it searches in the current directory for all the go test files
+// and then all the relevant function names.
+// for test names use prefix of 'Test' or 'Example', and for benchmark
+// test names use 'Benchmark'
+func predictTest(funcPrefix ...string) complete.Predicate {
 	return func(last string) []match.Matcher {
-		tests := testNames(testType)
+		tests := testNames(funcPrefix)
 		options := make([]match.Matcher, len(tests))
 		for i := range tests {
 			options[i] = match.Prefix(tests[i])
@@ -24,20 +29,20 @@ func predictTest(testType string) complete.Predicate {
 }
 
 // get all test names in current directory
-func testNames(testType string) (tests []string) {
+func testNames(funcPrefix []string) (tests []string) {
 	filepath.Walk("./", func(path string, info os.FileInfo, err error) error {
 		// if not a test file, skip
 		if !strings.HasSuffix(path, "_test.go") {
 			return nil
 		}
 		// inspect test file and append all the test names
-		tests = append(tests, testsInFile(testType, path)...)
+		tests = append(tests, testsInFile(funcPrefix, path)...)
 		return nil
 	})
 	return
 }
 
-func testsInFile(testType, path string) (tests []string) {
+func testsInFile(funcPrefix []string, path string) (tests []string) {
 	fset := token.NewFileSet()
 	f, err := parser.ParseFile(fset, path, nil, 0)
 	if err != nil {
@@ -47,8 +52,11 @@ func testsInFile(testType, path string) (tests []string) {
 	for _, d := range f.Decls {
 		if f, ok := d.(*ast.FuncDecl); ok {
 			name := f.Name.String()
-			if strings.HasPrefix(name, testType) {
-				tests = append(tests, name)
+			for _, prefix := range funcPrefix {
+				if strings.HasPrefix(name, prefix) {
+					tests = append(tests, name)
+					break
+				}
 			}
 		}
 	}