internal/lsp: add `run test` code lens

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

internal/lsp/command.go

@@ -6,18 +6,43 @@ package lsp
 
 import (
 	"context"
+	"io"
+	"path/filepath"
 	"strings"
 
+	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/gocommand"
+	"golang.org/x/tools/internal/lsp/debug/tag"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/packagesinternal"
 	"golang.org/x/tools/internal/xcontext"
+	"golang.org/x/xerrors"
 	errors "golang.org/x/xerrors"
 )
 
 func (s *Server) executeCommand(ctx context.Context, params *protocol.ExecuteCommandParams) (interface{}, error) {
 	switch params.Command {
+	case source.CommandTest:
+		if len(s.session.UnsavedFiles()) != 0 {
+			return nil, s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
+				Type:    protocol.Error,
+				Message: "could not run tests, there are unsaved files in the view",
+			})
+		}
+
+		funcName, uri, err := getRunTestArguments(params.Arguments)
+		if err != nil {
+			return nil, err
+		}
+
+		snapshot, fh, ok, err := s.beginFileRequest(protocol.DocumentURI(uri), source.Go)
+		if !ok {
+			return nil, err
+		}
+
+		dir := filepath.Dir(fh.Identity().URI.Filename())
+		go s.runTest(ctx, funcName, dir, snapshot)
 	case source.CommandGenerate:
 		dir, recursive, err := getGenerateRequest(params.Arguments)
 		if err != nil {
@@ -71,6 +96,57 @@ func (s *Server) executeCommand(ctx context.Context, params *protocol.ExecuteCom
 	return nil, nil
 }
 
+func (s *Server) runTest(ctx context.Context, funcName string, dir string, snapshot source.Snapshot) {
+	args := []string{"-run", funcName, dir}
+	inv := gocommand.Invocation{
+		Verb:       "test",
+		Args:       args,
+		Env:        snapshot.Config(ctx).Env,
+		WorkingDir: dir,
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	er := &eventWriter{ctx: ctx, operation: "test"}
+	wc := s.newProgressWriter(ctx, "test", "running "+funcName, cancel)
+	defer wc.Close()
+
+	messageType := protocol.Info
+	message := "test passed"
+	stderr := io.MultiWriter(er, wc)
+	if err := inv.RunPiped(ctx, er, stderr); err != nil {
+		event.Error(ctx, "test: command error", err, tag.Directory.Of(dir))
+		if !xerrors.Is(err, context.Canceled) {
+			messageType = protocol.Error
+			message = "test failed"
+		}
+	}
+
+	s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
+		Type:    messageType,
+		Message: message,
+	})
+}
+
+func getRunTestArguments(args []interface{}) (string, string, error) {
+	if len(args) != 2 {
+		return "", "", errors.Errorf("expected one test func name and one file path, got %v", args)
+	}
+
+	funcName, ok := args[0].(string)
+	if !ok {
+		return "", "", errors.Errorf("expected func name to be a string, got %T", args[0])
+	}
+
+	file, ok := args[1].(string)
+	if !ok {
+		return "", "", errors.Errorf("expected file to be a string, got %T", args[1])
+	}
+
+	return funcName, file, nil
+}
+
 func getGenerateRequest(args []interface{}) (string, bool, error) {
 	if len(args) != 2 {
 		return "", false, errors.Errorf("expected exactly 2 arguments but got %d", len(args))

internal/lsp/generate.go

@@ -23,8 +23,8 @@ func (s *Server) runGenerate(ctx context.Context, dir string, recursive bool) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	er := &eventWriter{ctx: ctx}
-	wc := s.newProgressWriter(ctx, cancel)
+	er := &eventWriter{ctx: ctx, operation: "generate"}
+	wc := s.newProgressWriter(ctx, GenerateWorkDoneTitle, "running go generate", cancel)
 	defer wc.Close()
 	args := []string{"-x"}
 	if recursive {
@@ -53,20 +53,20 @@ func (s *Server) runGenerate(ctx context.Context, dir string, recursive bool) {
 // event.Print with the operation=generate tag
 // to distinguish its logs from others.
 type eventWriter struct {
-	ctx context.Context
+	ctx       context.Context
+	operation string
 }
 
 func (ew *eventWriter) Write(p []byte) (n int, err error) {
-	event.Log(ew.ctx, string(p), tag.Operation.Of("generate"))
+	event.Log(ew.ctx, string(p), tag.Operation.Of(ew.operation))
 	return len(p), nil
 }
 
 // newProgressWriter returns an io.WriterCloser that can be used
-// to report progress on the "go generate" command based on the
-// client capabilities.
-func (s *Server) newProgressWriter(ctx context.Context, cancel func()) io.WriteCloser {
+// to report progress on a command based on the client capabilities.
+func (s *Server) newProgressWriter(ctx context.Context, title, message string, cancel func()) io.WriteCloser {
 	if s.supportsWorkDoneProgress {
-		wd := s.StartWork(ctx, GenerateWorkDoneTitle, "running go generate", cancel)
+		wd := s.StartWork(ctx, title, message, cancel)
 		return &workDoneWriter{ctx, wd}
 	}
 	mw := &messageWriter{ctx, cancel, s.client}

internal/lsp/source/code_lens.go

@@ -6,8 +6,11 @@ package source
 
 import (
 	"context"
+	"go/ast"
 	"go/token"
+	"go/types"
 	"path/filepath"
+	"regexp"
 	"strings"
 
 	"golang.org/x/tools/internal/lsp/protocol"
@@ -15,13 +18,112 @@ import (
 
 // CodeLens computes code lens for Go source code.
 func CodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]protocol.CodeLens, error) {
-	if !snapshot.View().Options().EnabledCodeLens[CommandGenerate] {
-		return nil, nil
-	}
 	f, _, m, _, err := snapshot.View().Session().Cache().ParseGoHandle(fh, ParseFull).Parse(ctx)
 	if err != nil {
 		return nil, err
 	}
+
+	var codeLens []protocol.CodeLens
+
+	if snapshot.View().Options().EnabledCodeLens[CommandGenerate] {
+		ggcl, err := goGenerateCodeLens(ctx, snapshot, fh, f, m)
+		if err != nil {
+			return nil, err
+		}
+		codeLens = append(codeLens, ggcl...)
+	}
+
+	if snapshot.View().Options().EnabledCodeLens[CommandTest] {
+		rtcl, err := runTestCodeLens(ctx, snapshot, fh, f, m)
+		if err != nil {
+			return nil, err
+		}
+		codeLens = append(codeLens, rtcl...)
+	}
+
+	return codeLens, nil
+}
+
+var testMatcher = regexp.MustCompile("^Test[^a-z]")
+var benchMatcher = regexp.MustCompile("^Benchmark[^a-z]")
+
+func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle, f *ast.File, m *protocol.ColumnMapper) ([]protocol.CodeLens, error) {
+	codeLens := make([]protocol.CodeLens, 0)
+
+	pkg, _, err := getParsedFile(ctx, snapshot, fh, WidestPackageHandle)
+	if err != nil {
+		return nil, err
+	}
+
+	if !strings.HasSuffix(fh.Identity().URI.Filename(), "_test.go") {
+		return nil, nil
+	}
+
+	for _, d := range f.Decls {
+		fn, ok := d.(*ast.FuncDecl)
+		if !ok {
+			continue
+		}
+
+		if isTestFunc(fn, pkg) {
+			fset := snapshot.View().Session().Cache().FileSet()
+			rng, err := newMappedRange(fset, m, d.Pos(), d.Pos()).Range()
+			if err != nil {
+				return nil, err
+			}
+
+			uri := fh.Identity().URI
+			codeLens = append(codeLens, protocol.CodeLens{
+				Range: rng,
+				Command: protocol.Command{
+					Title:     "run test",
+					Command:   "test",
+					Arguments: []interface{}{fn.Name.Name, uri},
+				},
+			})
+		}
+	}
+
+	return codeLens, nil
+}
+
+func isTestFunc(fn *ast.FuncDecl, pkg Package) bool {
+	typesInfo := pkg.GetTypesInfo()
+	if typesInfo == nil {
+		return false
+	}
+
+	sig, ok := typesInfo.ObjectOf(fn.Name).Type().(*types.Signature)
+	if !ok {
+		return false
+	}
+
+	// test funcs should have a single parameter, so we can exit early if that's not the case.
+	if sig.Params().Len() != 1 {
+		return false
+	}
+
+	firstParam, ok := sig.Params().At(0).Type().(*types.Pointer)
+	if !ok {
+		return false
+	}
+
+	firstParamElem, ok := firstParam.Elem().(*types.Named)
+	if !ok {
+		return false
+	}
+
+	firstParamObj := firstParamElem.Obj()
+	if firstParamObj.Pkg().Path() != "testing" {
+		return false
+	}
+
+	firstParamName := firstParamObj.Id()
+	return (firstParamName == "T" && testMatcher.MatchString(fn.Name.Name)) ||
+		(firstParamName == "B" && benchMatcher.MatchString(fn.Name.Name))
+}
+
+func goGenerateCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle, f *ast.File, m *protocol.ColumnMapper) ([]protocol.CodeLens, error) {
 	const ggDirective = "//go:generate"
 	for _, c := range f.Comments {
 		for _, l := range c.List {

internal/lsp/source/options.go

@@ -53,6 +53,8 @@ import (
 )
 
 const (
+	// CommandGenerate is a gopls command to run `go test` for a specific test function.
+	CommandTest = "test"
 	// CommandGenerate is a gopls command to run `go generate` for a directory.
 	CommandGenerate = "generate"
 	// CommandTidy is a gopls command to run `go mod tidy` for a module.
@@ -88,6 +90,7 @@ func DefaultOptions() Options {
 				Sum: {},
 			},
 			SupportedCommands: []string{
+				CommandTest,              // for "go test" commands
 				CommandTidy,              // for go.mod files
 				CommandUpgradeDependency, // for go.mod dependency upgrades
 				CommandGenerate,          // for "go generate" commands

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

@@ -0,0 +1,16 @@
+package codelens
+
+import "testing"
+
+// no code lens for TestMain
+func TestMain(m *testing.M) {
+}
+
+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")
+}

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

@@ -1,5 +1,5 @@
 -- summary --
-CodeLensCount = 2
+CodeLensCount = 4
 CompletionsCount = 240
 CompletionSnippetCount = 76
 UnimportedCompletionsCount = 6

internal/lsp/tests/tests.go

@@ -222,6 +222,7 @@ func DefaultOptions() source.Options {
 		},
 		source.Sum: {},
 	}
+	o.UserOptions.EnabledCodeLens[source.CommandTest] = true
 	o.HoverKind = source.SynopsisDocumentation
 	o.InsertTextFormat = protocol.SnippetTextFormat
 	o.CompletionBudget = time.Minute