gopls: add "go test" code action

This change adds "go test" as a code action and also introduce the
concept of explicit code actions, i.e. code actions that aren't returned
to the client unless it explicitly ask for it.

The purpose is to be able to have a mechanism that allows users to
execute a specific command in one shot, and future CL:s will add more of
the existing code lenses as explicit code actions. Code lenses can't be
used directly since they lack the range/kind combo to make them unique.

Updates golang/go#40438

Change-Id: I245df4e26c9e02e529662181e2c1bc44f999e101
Reviewed-on: https://go-review.googlesource.com/c/tools/+/259377
Trust: Rebecca Stambler <rstambler@golang.org>
Trust: Peter Weinberger <pjw@google.com>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Go Bot <gobot@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/code_action.go b/internal/lsp/code_action.go
index 090ffd1..20e9f99 100644
--- a/internal/lsp/code_action.go
+++ b/internal/lsp/code_action.go
@@ -34,14 +34,22 @@
 	}
 
 	// The Only field of the context specifies which code actions the client wants.
-	// If Only is empty, assume that the client wants all of the possible code actions.
+	// If Only is empty, assume that the client wants all of the non-explicit code actions.
 	var wanted map[protocol.CodeActionKind]bool
+
+	// Explicit Code Actions are opt-in and shouldn't be returned to the client unless
+	// requested using Only.
+	// TODO: Add other CodeLenses such as GoGenerate, RegenerateCgo, etc..
+	explicit := map[protocol.CodeActionKind]bool{
+		protocol.GoTest: true,
+	}
+
 	if len(params.Context.Only) == 0 {
 		wanted = supportedCodeActions
 	} else {
 		wanted = make(map[protocol.CodeActionKind]bool)
 		for _, only := range params.Context.Only {
-			wanted[only] = supportedCodeActions[only]
+			wanted[only] = supportedCodeActions[only] || explicit[only]
 		}
 	}
 	if len(wanted) == 0 {
@@ -169,6 +177,15 @@
 			}
 			codeActions = append(codeActions, fixes...)
 		}
+
+		if wanted[protocol.GoTest] {
+			fixes, err := goTest(ctx, snapshot, uri, params.Range)
+			if err != nil {
+				return nil, err
+			}
+			codeActions = append(codeActions, fixes...)
+		}
+
 	default:
 		// Unsupported file kind for a code action.
 		return nil, nil
@@ -547,3 +564,46 @@
 		},
 	}, err
 }
+
+func goTest(ctx context.Context, snapshot source.Snapshot, uri span.URI, rng protocol.Range) ([]protocol.CodeAction, error) {
+	fh, err := snapshot.GetFile(ctx, uri)
+	if err != nil {
+		return nil, err
+	}
+	fns, err := source.TestsAndBenchmarks(ctx, snapshot, fh)
+	if err != nil {
+		return nil, err
+	}
+
+	var tests, benchmarks []string
+	for _, fn := range fns.Tests {
+		if !protocol.Intersect(fn.Rng, rng) {
+			continue
+		}
+		tests = append(tests, fn.Name)
+	}
+	for _, fn := range fns.Benchmarks {
+		if !protocol.Intersect(fn.Rng, rng) {
+			continue
+		}
+		benchmarks = append(benchmarks, fn.Name)
+	}
+
+	if len(tests) == 0 && len(benchmarks) == 0 {
+		return nil, nil
+	}
+
+	jsonArgs, err := source.MarshalArgs(uri, tests, benchmarks)
+	if err != nil {
+		return nil, err
+	}
+	return []protocol.CodeAction{{
+		Title: source.CommandTest.Name,
+		Kind:  protocol.GoTest,
+		Command: &protocol.Command{
+			Title:     source.CommandTest.Title,
+			Command:   source.CommandTest.ID(),
+			Arguments: jsonArgs,
+		},
+	}}, nil
+}
diff --git a/internal/lsp/protocol/codeactionkind.go b/internal/lsp/protocol/codeactionkind.go
new file mode 100644
index 0000000..9a95800
--- /dev/null
+++ b/internal/lsp/protocol/codeactionkind.go
@@ -0,0 +1,11 @@
+// Copyright 2020 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 protocol
+
+// Custom code actions that aren't explicitly stated in LSP
+const (
+	GoTest CodeActionKind = "goTest"
+	// TODO: Add GoGenerate, RegenerateCgo etc.
+)
diff --git a/internal/lsp/source/code_lens.go b/internal/lsp/source/code_lens.go
index 6bf3c57..d6d4491 100644
--- a/internal/lsp/source/code_lens.go
+++ b/internal/lsp/source/code_lens.go
@@ -37,64 +37,50 @@
 func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]protocol.CodeLens, error) {
 	codeLens := make([]protocol.CodeLens, 0)
 
-	if !strings.HasSuffix(fh.URI().Filename(), "_test.go") {
-		return nil, nil
-	}
-	pkg, pgf, err := GetParsedFile(ctx, snapshot, fh, WidestPackage)
+	fns, err := TestsAndBenchmarks(ctx, snapshot, fh)
 	if err != nil {
 		return nil, err
 	}
-
-	var benchFns []string
-	for _, d := range pgf.File.Decls {
-		fn, ok := d.(*ast.FuncDecl)
-		if !ok {
-			continue
-		}
-		if benchmarkRe.MatchString(fn.Name.Name) {
-			benchFns = append(benchFns, fn.Name.Name)
-		}
-		rng, err := NewMappedRange(snapshot.FileSet(), pgf.Mapper, d.Pos(), d.Pos()).Range()
+	for _, fn := range fns.Tests {
+		jsonArgs, err := MarshalArgs(fh.URI(), []string{fn.Name}, nil)
 		if err != nil {
 			return nil, err
 		}
+		codeLens = append(codeLens, protocol.CodeLens{
+			Range: protocol.Range{Start: fn.Rng.Start, End: fn.Rng.Start},
+			Command: protocol.Command{
+				Title:     "run test",
+				Command:   CommandTest.ID(),
+				Arguments: jsonArgs,
+			},
+		})
+	}
 
-		if matchTestFunc(fn, pkg, testRe, "T") {
-			jsonArgs, err := MarshalArgs(fh.URI(), []string{fn.Name.Name}, nil)
-			if err != nil {
-				return nil, err
-			}
-			codeLens = append(codeLens, protocol.CodeLens{
-				Range: rng,
-				Command: protocol.Command{
-					Title:     "run test",
-					Command:   CommandTest.ID(),
-					Arguments: jsonArgs,
-				},
-			})
+	for _, fn := range fns.Benchmarks {
+		jsonArgs, err := MarshalArgs(fh.URI(), nil, []string{fn.Name})
+		if err != nil {
+			return nil, err
 		}
+		codeLens = append(codeLens, protocol.CodeLens{
+			Range: protocol.Range{Start: fn.Rng.Start, End: fn.Rng.Start},
+			Command: protocol.Command{
+				Title:     "run benchmark",
+				Command:   CommandTest.ID(),
+				Arguments: jsonArgs,
+			},
+		})
+	}
 
-		if matchTestFunc(fn, pkg, benchmarkRe, "B") {
-			jsonArgs, err := MarshalArgs(fh.URI(), nil, []string{fn.Name.Name})
-			if err != nil {
-				return nil, err
-			}
-			codeLens = append(codeLens, protocol.CodeLens{
-				Range: rng,
-				Command: protocol.Command{
-					Title:     "run benchmark",
-					Command:   CommandTest.ID(),
-					Arguments: jsonArgs,
-				},
-			})
-		}
+	_, pgf, err := GetParsedFile(ctx, snapshot, fh, WidestPackage)
+	if err != nil {
+		return nil, err
 	}
 	// add a code lens to the top of the file which runs all benchmarks in the file
 	rng, err := NewMappedRange(snapshot.FileSet(), pgf.Mapper, pgf.File.Package, pgf.File.Package).Range()
 	if err != nil {
 		return nil, err
 	}
-	args, err := MarshalArgs(fh.URI(), []string{}, benchFns)
+	args, err := MarshalArgs(fh.URI(), []string{}, fns.Benchmarks)
 	if err != nil {
 		return nil, err
 	}
@@ -109,6 +95,50 @@
 	return codeLens, nil
 }
 
+type testFn struct {
+	Name string
+	Rng  protocol.Range
+}
+
+type testFns struct {
+	Tests      []testFn
+	Benchmarks []testFn
+}
+
+func TestsAndBenchmarks(ctx context.Context, snapshot Snapshot, fh FileHandle) (testFns, error) {
+	var out testFns
+
+	if !strings.HasSuffix(fh.URI().Filename(), "_test.go") {
+		return out, nil
+	}
+	pkg, pgf, err := GetParsedFile(ctx, snapshot, fh, WidestPackage)
+	if err != nil {
+		return out, err
+	}
+
+	for _, d := range pgf.File.Decls {
+		fn, ok := d.(*ast.FuncDecl)
+		if !ok {
+			continue
+		}
+
+		rng, err := NewMappedRange(snapshot.FileSet(), pgf.Mapper, d.Pos(), fn.End()).Range()
+		if err != nil {
+			return out, err
+		}
+
+		if matchTestFunc(fn, pkg, testRe, "T") {
+			out.Tests = append(out.Tests, testFn{fn.Name.Name, rng})
+		}
+
+		if matchTestFunc(fn, pkg, benchmarkRe, "B") {
+			out.Benchmarks = append(out.Benchmarks, testFn{fn.Name.Name, rng})
+		}
+	}
+
+	return out, nil
+}
+
 func matchTestFunc(fn *ast.FuncDecl, pkg Package, nameRe *regexp.Regexp, paramID string) bool {
 	// Make sure that the function name matches a test function.
 	if !nameRe.MatchString(fn.Name.Name) {