internal/lsp: separate test and benchmark codelens

This CL adds benchmarking by using the "-bench" flag for the test
command.

Change-Id: Idf714de1c6c3350d26a6874e7b4278927d0336fd
Reviewed-on: https://go-review.googlesource.com/c/tools/+/243159
Run-TryBot: Robert Findley <rfindley@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>

internal/lsp/command.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"strings"
 
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/lsp/debug/tag"
@@ -50,15 +51,16 @@ func (s *Server) executeCommand(ctx context.Context, params *protocol.ExecuteCom
 	switch params.Command {
 	case source.CommandTest:
 		var uri protocol.DocumentURI
+		var flag string
 		var funcName string
-		if err := source.DecodeArgs(params.Arguments, &uri, &funcName); err != nil {
+		if err := source.DecodeArgs(params.Arguments, &uri, &flag, &funcName); err != nil {
 			return nil, err
 		}
 		snapshot, _, ok, err := s.beginFileRequest(ctx, uri, source.UnknownKind)
 		if !ok {
 			return nil, err
 		}
-		go s.runGoTest(ctx, snapshot, funcName)
+		go s.runTest(ctx, snapshot, []string{flag, funcName})
 	case source.CommandGenerate:
 		var uri protocol.DocumentURI
 		var recursive bool
@@ -140,12 +142,12 @@ func (s *Server) directGoModCommand(ctx context.Context, uri protocol.DocumentUR
 	return view.Snapshot().RunGoCommandDirect(ctx, verb, args)
 }
 
-func (s *Server) runGoTest(ctx context.Context, snapshot source.Snapshot, funcName string) error {
+func (s *Server) runTest(ctx context.Context, snapshot source.Snapshot, args []string) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	ew := &eventWriter{ctx: ctx, operation: "test"}
-	msg := fmt.Sprintf("testing %s", funcName)
+	msg := fmt.Sprintf("running `go test %s`", strings.Join(args, " "))
 	wc := s.newProgressWriter(ctx, "test", msg, msg, cancel)
 	defer wc.Close()
 
@@ -153,7 +155,7 @@ func (s *Server) runGoTest(ctx context.Context, snapshot source.Snapshot, funcNa
 	message := "test passed"
 	stderr := io.MultiWriter(ew, wc)
 
-	if err := snapshot.RunGoCommandPiped(ctx, "test", []string{"-run", funcName}, ew, stderr); err != nil {
+	if err := snapshot.RunGoCommandPiped(ctx, "test", args, ew, stderr); err != nil {
 		if errors.Is(err, context.Canceled) {
 			return err
 		}

internal/lsp/source/code_lens.go

@@ -41,6 +41,9 @@ func CodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]protocol
 	return result, nil
 }
 
+var testRe = regexp.MustCompile("^Test[^a-z]")
+var benchmarkRe = regexp.MustCompile("^Benchmark[^a-z]")
+
 func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]protocol.CodeLens, error) {
 	codeLens := make([]protocol.CodeLens, 0)
 
@@ -60,38 +63,48 @@ func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]p
 		if !ok {
 			continue
 		}
-		if !isTestFunc(fn, pkg) {
-			continue
-		}
 		fset := snapshot.View().Session().Cache().FileSet()
 		rng, err := newMappedRange(fset, m, d.Pos(), d.Pos()).Range()
 		if err != nil {
 			return nil, err
 		}
-		jsonArgs, err := EncodeArgs(fh.URI(), fn.Name.Name)
-		if err != nil {
-			return nil, err
+
+		if matchTestFunc(fn, pkg, testRe, "T") {
+			jsonArgs, err := EncodeArgs(fh.URI(), "-run", fn.Name.Name)
+			if err != nil {
+				return nil, err
+			}
+			codeLens = append(codeLens, protocol.CodeLens{
+				Range: rng,
+				Command: protocol.Command{
+					Title:     "run test",
+					Command:   CommandTest,
+					Arguments: jsonArgs,
+				},
+			})
+		}
+
+		if matchTestFunc(fn, pkg, benchmarkRe, "B") {
+			jsonArgs, err := EncodeArgs(fh.URI(), "-bench", fn.Name.Name)
+			if err != nil {
+				return nil, err
+			}
+			codeLens = append(codeLens, protocol.CodeLens{
+				Range: rng,
+				Command: protocol.Command{
+					Title:     "run benchmark",
+					Command:   CommandTest,
+					Arguments: jsonArgs,
+				},
+			})
 		}
-		codeLens = append(codeLens, protocol.CodeLens{
-			Range: rng,
-			Command: protocol.Command{
-				Title:     "run test",
-				Command:   CommandTest,
-				Arguments: jsonArgs,
-			},
-		})
 	}
 	return codeLens, nil
 }
 
-var (
-	testRe      = regexp.MustCompile("^Test[^a-z]")
-	benchmarkRe = regexp.MustCompile("^Benchmark[^a-z]")
-)
-
-func isTestFunc(fn *ast.FuncDecl, pkg Package) bool {
-	// Make sure that the function name matches either a test or benchmark function.
-	if !(testRe.MatchString(fn.Name.Name) || benchmarkRe.MatchString(fn.Name.Name)) {
+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) {
 		return false
 	}
 	info := pkg.GetTypesInfo()
@@ -111,8 +124,7 @@ func isTestFunc(fn *ast.FuncDecl, pkg Package) bool {
 		return false
 	}
 
-	// Check the type of the only parameter to confirm that it is *testing.T
-	// or *testing.B.
+	// Check the type of the only parameter
 	paramTyp, ok := sig.Params().At(0).Type().(*types.Pointer)
 	if !ok {
 		return false
@@ -125,8 +137,7 @@ func isTestFunc(fn *ast.FuncDecl, pkg Package) bool {
 	if namedObj.Pkg().Path() != "testing" {
 		return false
 	}
-	paramName := namedObj.Id()
-	return paramName == "T" || paramName == "B"
+	return namedObj.Id() == paramID
 }
 
 func goGenerateCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]protocol.CodeLens, error) {

internal/lsp/testdata/lsp/primarymod/codelens/codelens_test.go

@@ -10,7 +10,7 @@ func TestFuncWithCodeLens(t *testing.T) { //@codelens("func", "run test", "test"
 func thisShouldNotHaveACodeLens(t *testing.T) {
 }
 
-func BenchmarkFuncWithCodeLens(b *testing.B) { //@codelens("func", "run test", "test")
+func BenchmarkFuncWithCodeLens(b *testing.B) { //@codelens("func", "run benchmark", "test")
 }
 
 func helper() {} // expect no code lens