internal/lsp: Provide completions for test function definitions

In test files, function definitions starting with
Test, Bench, or Fuzz can be completed almost automatically.

For the snippets the user hits tab, completes the name, hits
tab again, and the function is defined, except (of course) for its
body.

Otherwise a completion that fills in the signature is proposed.

Where appropriate, 'TestMain(m *testing.M)' is also offered as
a completion.

Fixes golang/go#46896 and golang/go#51089

Change-Id: I46c05af0ead79c1d82ca40b2c605045e06e1a35d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/385974
Run-TryBot: Peter Weinberger <[email protected]>
Trust: Peter Weinberger <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
gopls-CI: kokoro <[email protected]>
Reviewed-by: Hyang-Ah Hana Kim <[email protected]>

Author: pjweinbgo

gopls/internal/regtest/completion/completion_test.go

@@ -256,7 +256,7 @@ func compareCompletionResults(want []string, gotItems []protocol.CompletionItem)
 
 	for i, v := range got {
 		if v != want[i] {
-			return fmt.Sprintf("completion results are not the same: got %v, want %v", got, want)
+			return fmt.Sprintf("%d completion result not the same: got %q, want %q", i, v, want[i])
 		}
 	}
 
@@ -546,3 +546,56 @@ func main() {
 		}
 	})
 }
+
+func TestDefinition(t *testing.T) {
+	stuff := `
+-- go.mod --
+module mod.com
+
+go 1.18
+-- a_test.go --
+package foo
+func T()
+func TestG()
+func TestM()
+func TestMi()
+func Ben()
+func Fuz()
+func Testx()
+func TestMe(t *testing.T)
+func BenchmarkFoo()
+`
+	// All those parentheses are needed for the completion code to see
+	// later lines as being definitions
+	tests := []struct {
+		pat  string
+		want []string
+	}{
+		{"T", []string{"TestXxx(t *testing.T)", "TestMain(m *testing.M)"}},
+		{"TestM", []string{"TestMain(m *testing.M)", "TestM(t *testing.T)"}},
+		{"TestMi", []string{"TestMi(t *testing.T)"}},
+		{"TestG", []string{"TestG(t *testing.T)"}},
+		{"B", []string{"BenchmarkXxx(b *testing.B)"}},
+		{"BenchmarkFoo", []string{"BenchmarkFoo(b *testing.B)"}},
+		{"F", []string{"FuzzXxx(f *testing.F)"}},
+		{"Testx", nil},
+		{"TestMe", []string{"TestMe"}},
+	}
+	fname := "a_test.go"
+	Run(t, stuff, func(t *testing.T, env *Env) {
+		env.OpenFile(fname)
+		env.Await(env.DoneWithOpen())
+		for _, tst := range tests {
+			pos := env.RegexpSearch(fname, tst.pat)
+			pos.Column += len(tst.pat)
+			completions := env.Completion(fname, pos)
+			result := compareCompletionResults(tst.want, completions.Items)
+			if result != "" {
+				t.Errorf("%s failed: %s:%q", tst.pat, result, tst.want)
+				for i, it := range completions.Items {
+					t.Errorf("%d got %q %q", i, it.Label, it.Detail)
+				}
+			}
+		}
+	})
+}

internal/lsp/source/completion/completion.go

@@ -485,6 +485,13 @@ func Completion(ctx context.Context, snapshot source.Snapshot, fh source.FileHan
 					qual := types.RelativeTo(pkg.GetTypes())
 					objStr = types.ObjectString(obj, qual)
 				}
+				ans, sel := definition(path, obj, snapshot.FileSet(), pgf.Mapper, fh)
+				if ans != nil {
+					sort.Slice(ans, func(i, j int) bool {
+						return ans[i].Score > ans[j].Score
+					})
+					return ans, sel, nil
+				}
 				return nil, nil, ErrIsDefinition{objStr: objStr}
 			}
 		}

internal/lsp/source/completion/definition.go

@@ -0,0 +1,127 @@
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package completion
+
+import (
+	"go/ast"
+	"go/token"
+	"go/types"
+	"strings"
+	"unicode"
+	"unicode/utf8"
+
+	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/lsp/snippet"
+	"golang.org/x/tools/internal/lsp/source"
+)
+
+// some definitions can be completed
+// So far, TestFoo(t *testing.T), TestMain(m *testing.M)
+// BenchmarkFoo(b *testing.B), FuzzFoo(f *testing.F)
+
+// path[0] is known to be *ast.Ident
+func definition(path []ast.Node, obj types.Object, fset *token.FileSet, mapper *protocol.ColumnMapper, fh source.FileHandle) ([]CompletionItem, *Selection) {
+	if _, ok := obj.(*types.Func); !ok {
+		return nil, nil // not a function at all
+	}
+	if !strings.HasSuffix(fh.URI().Filename(), "_test.go") {
+		return nil, nil
+	}
+
+	name := path[0].(*ast.Ident).Name
+	if len(name) == 0 {
+		// can't happen
+		return nil, nil
+	}
+	pos := path[0].Pos()
+	sel := &Selection{
+		content:     "",
+		cursor:      pos,
+		MappedRange: source.NewMappedRange(fset, mapper, pos, pos),
+	}
+	var ans []CompletionItem
+
+	// Always suggest TestMain, if possible
+	if strings.HasPrefix("TestMain", name) {
+		ans = []CompletionItem{defItem("TestMain(m *testing.M)", obj)}
+	}
+
+	// If a snippet is possible, suggest it
+	if strings.HasPrefix("Test", name) {
+		ans = append(ans, defSnippet("Test", "Xxx", "(t *testing.T)", obj))
+		return ans, sel
+	} else if strings.HasPrefix("Benchmark", name) {
+		ans = append(ans, defSnippet("Benchmark", "Xxx", "(b *testing.B)", obj))
+		return ans, sel
+	} else if strings.HasPrefix("Fuzz", name) {
+		ans = append(ans, defSnippet("Fuzz", "Xxx", "(f *testing.F)", obj))
+		return ans, sel
+	}
+
+	// Fill in the argument for what the user has already typed
+	if got := defMatches(name, "Test", path, "(t *testing.T)"); got != "" {
+		ans = append(ans, defItem(got, obj))
+	} else if got := defMatches(name, "Benchmark", path, "(b *testing.B)"); got != "" {
+		ans = append(ans, defItem(got, obj))
+	} else if got := defMatches(name, "Fuzz", path, "(f *testing.F)"); got != "" {
+		ans = append(ans, defItem(got, obj))
+	}
+	return ans, sel
+}
+
+func defMatches(name, pat string, path []ast.Node, arg string) string {
+	idx := strings.Index(name, pat)
+	if idx < 0 {
+		return ""
+	}
+	c, _ := utf8.DecodeRuneInString(name[len(pat):])
+	if unicode.IsLower(c) {
+		return ""
+	}
+	fd, ok := path[1].(*ast.FuncDecl)
+	if !ok {
+		// we don't know what's going on
+		return ""
+	}
+	fp := fd.Type.Params
+	if fp != nil && len(fp.List) > 0 {
+		// signature already there, minimal suggestion
+		return name
+	}
+	// suggesting signature too
+	return name + arg
+}
+
+func defSnippet(prefix, placeholder, suffix string, obj types.Object) CompletionItem {
+	var sn snippet.Builder
+	sn.WriteText(prefix)
+	if placeholder != "" {
+		sn.WritePlaceholder(func(b *snippet.Builder) { b.WriteText(placeholder) })
+	}
+	sn.WriteText(suffix + " {\n")
+	sn.WriteFinalTabstop()
+	sn.WriteText("\n}")
+	return CompletionItem{
+		Label:         prefix + placeholder + suffix,
+		Detail:        "tab, type the rest of the name, then tab",
+		Kind:          protocol.FunctionCompletion,
+		Depth:         0,
+		Score:         10,
+		snippet:       &sn,
+		Documentation: prefix + " test function",
+		obj:           obj,
+	}
+}
+func defItem(val string, obj types.Object) CompletionItem {
+	return CompletionItem{
+		Label:         val,
+		InsertText:    val,
+		Kind:          protocol.FunctionCompletion,
+		Depth:         0,
+		Score:         9, // prefer the snippets when available
+		Documentation: "complete the parameter",
+		obj:           obj,
+	}
+}