internal/lsp: add "run file benchmarks" code lens

This CL adds a code lens to run all benchmarks in a file. Additionally,
it updates the test command handler to better support both tests and
benchmarks.

Updates golang/go#36787

Change-Id: I6e90460f7d97607f96c263be0754537764bd0052
Reviewed-on: https://go-review.googlesource.com/c/tools/+/246017
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

@@ -9,7 +9,6 @@ import (
 	"fmt"
 	"io"
 	"path"
-	"strings"
 
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/lsp/debug/tag"
@@ -97,9 +96,8 @@ func (s *Server) executeCommand(ctx context.Context, params *protocol.ExecuteCom
 	switch command {
 	case source.CommandTest:
 		var uri protocol.DocumentURI
-		var flag string
-		var funcName string
-		if err := source.UnmarshalArgs(params.Arguments, &uri, &flag, &funcName); err != nil {
+		var tests, benchmarks []string
+		if err := source.UnmarshalArgs(params.Arguments, &uri, &tests, &benchmarks); err != nil {
 			return nil, err
 		}
 		snapshot, _, ok, release, err := s.beginFileRequest(ctx, uri, source.UnknownKind)
@@ -107,7 +105,7 @@ func (s *Server) executeCommand(ctx context.Context, params *protocol.ExecuteCom
 		if !ok {
 			return nil, err
 		}
-		go s.runTest(ctx, snapshot, []string{flag, funcName}, params.WorkDoneToken)
+		go s.runTests(ctx, snapshot, uri, params.WorkDoneToken, tests, benchmarks)
 	case source.CommandGenerate:
 		var uri protocol.DocumentURI
 		var recursive bool
@@ -193,26 +191,74 @@ func (s *Server) directGoModCommand(ctx context.Context, uri protocol.DocumentUR
 	return snapshot.RunGoCommandDirect(ctx, verb, args)
 }
 
-func (s *Server) runTest(ctx context.Context, snapshot source.Snapshot, args []string, token protocol.ProgressToken) error {
+func (s *Server) runTests(ctx context.Context, snapshot source.Snapshot, uri protocol.DocumentURI, token protocol.ProgressToken, tests, benchmarks []string) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
+	pkgs, err := snapshot.PackagesForFile(ctx, uri.SpanURI())
+	if err != nil {
+		return err
+	}
+	if len(pkgs) == 0 {
+		return fmt.Errorf("package could not be found for file: %s", uri.SpanURI().Filename())
+	}
+	pkgPath := pkgs[0].PkgPath()
+
+	// create output
 	ew := &eventWriter{ctx: ctx, operation: "test"}
-	msg := fmt.Sprintf("running `go test %s`", strings.Join(args, " "))
-	wc := s.progress.newWriter(ctx, "test", msg, msg, token, cancel)
+	var title string
+	if len(tests) > 0 && len(benchmarks) > 0 {
+		title = "tests and benchmarks"
+	} else if len(tests) > 0 {
+		title = "tests"
+	} else if len(benchmarks) > 0 {
+		title = "benchmarks"
+	} else {
+		return errors.New("No functions were provided")
+	}
+	msg := fmt.Sprintf("Running %s...", title)
+	wc := s.progress.newWriter(ctx, title, msg, msg, token, cancel)
 	defer wc.Close()
 
-	messageType := protocol.Info
-	message := "test passed"
 	stderr := io.MultiWriter(ew, wc)
 
-	if err := snapshot.RunGoCommandPiped(ctx, "test", args, ew, stderr); err != nil {
-		if errors.Is(err, context.Canceled) {
-			return err
+	// run `go test -run Func` on each test
+	var failedTests int
+	for _, funcName := range tests {
+		args := []string{pkgPath, "-run", fmt.Sprintf("^%s$", funcName)}
+		if err := snapshot.RunGoCommandPiped(ctx, "test", args, ew, stderr); err != nil {
+			if errors.Is(err, context.Canceled) {
+				return err
+			}
+			failedTests++
 		}
-		messageType = protocol.Error
-		message = "test failed"
 	}
+
+	// run `go test -run=^$ -bench Func` on each test
+	var failedBenchmarks int
+	for _, funcName := range tests {
+		args := []string{pkgPath, "-run=^$", "-bench", fmt.Sprintf("^%s$", funcName)}
+		if err := snapshot.RunGoCommandPiped(ctx, "test", args, ew, stderr); err != nil {
+			if errors.Is(err, context.Canceled) {
+				return err
+			}
+			failedBenchmarks++
+		}
+	}
+
+	messageType := protocol.Info
+	message := fmt.Sprintf("all %s passed", title)
+	if failedTests > 0 || failedBenchmarks > 0 {
+		messageType = protocol.Error
+	}
+	if failedTests > 0 && failedBenchmarks > 0 {
+		message = fmt.Sprintf("%d / %d tests failed and %d / %d benchmarks failed", failedTests, len(tests), failedBenchmarks, len(benchmarks))
+	} else if failedTests > 0 {
+		message = fmt.Sprintf("%d / %d tests failed", failedTests, len(tests))
+	} else if failedBenchmarks > 0 {
+		message = fmt.Sprintf("%d / %d benchmarks failed", failedBenchmarks, len(benchmarks))
+	}
+
 	return s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
 		Type:    messageType,
 		Message: message,

internal/lsp/source/code_lens.go

@@ -57,18 +57,23 @@ func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]p
 	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()
 		if err != nil {
 			return nil, err
 		}
 
 		if matchTestFunc(fn, pkg, testRe, "T") {
-			jsonArgs, err := MarshalArgs(fh.URI(), "-run", fn.Name.Name)
+			jsonArgs, err := MarshalArgs(fh.URI(), []string{fn.Name.Name}, nil)
 			if err != nil {
 				return nil, err
 			}
@@ -83,7 +88,7 @@ func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]p
 		}
 
 		if matchTestFunc(fn, pkg, benchmarkRe, "B") {
-			jsonArgs, err := MarshalArgs(fh.URI(), "-bench", fn.Name.Name)
+			jsonArgs, err := MarshalArgs(fh.URI(), nil, []string{fn.Name.Name})
 			if err != nil {
 				return nil, err
 			}
@@ -97,6 +102,23 @@ func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]p
 			})
 		}
 	}
+	// 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)
+	if err != nil {
+		return nil, err
+	}
+	codeLens = append(codeLens, protocol.CodeLens{
+		Range: rng,
+		Command: protocol.Command{
+			Title:     "run file benchmarks",
+			Command:   CommandTest.Name,
+			Arguments: args,
+		},
+	})
 	return codeLens, nil
 }
 

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

@@ -1,4 +1,4 @@
-package codelens
+package codelens //@codelens("package codelens", "run file benchmarks", "test")
 
 import "testing"
 

internal/lsp/testdata/lsp/summary.txt.golden

@@ -1,6 +1,6 @@
 -- summary --
 CallHierarchyCount = 1
-CodeLensCount = 4
+CodeLensCount = 5
 CompletionsCount = 239
 CompletionSnippetCount = 81
 UnimportedCompletionsCount = 6