gopls/internal/golang: add source code action for add test

This CL is some glue code which build the connection between the
LSP "code action request" with second call which compute the actual
DocumentChange.

AddTest source code action will create a test file if not already
exist and insert a random function at the end of the test file.

For testing, an internal boolean option "addTestSourceCodeAction"
is created and only effective if set explicitly in marker test.

For golang/vscode-go#1594

Change-Id: Ie3d9279ea2858805254181608a0d5103afd3a4c6
Reviewed-on: https://go-review.googlesource.com/c/tools/+/621056
Reviewed-by: Robert Findley <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

gopls/internal/cache/parsego/file.go

@@ -89,6 +89,11 @@ func (pgf *File) NodeRange(node ast.Node) (protocol.Range, error) {
 	return pgf.Mapper.NodeRange(pgf.Tok, node)
 }
 
+// NodeOffsets returns offsets for the ast.Node.
+func (pgf *File) NodeOffsets(node ast.Node) (start int, end int, _ error) {
+	return safetoken.Offsets(pgf.Tok, node.Pos(), node.End())
+}
+
 // NodeMappedRange returns a MappedRange for the ast.Node interval in this file.
 // A MappedRange can be converted to any other form.
 func (pgf *File) NodeMappedRange(node ast.Node) (protocol.MappedRange, error) {

gopls/internal/golang/addtest.go

@@ -0,0 +1,117 @@
+// Copyright 2019 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 golang
+
+// This file defines the behavior of the "Add test for FUNC" command.
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"go/token"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"golang.org/x/tools/gopls/internal/cache"
+	"golang.org/x/tools/gopls/internal/cache/parsego"
+	"golang.org/x/tools/gopls/internal/protocol"
+)
+
+// AddTestForFunc adds a test for the function enclosing the given input range.
+// It creates a _test.go file if one does not already exist.
+func AddTestForFunc(ctx context.Context, snapshot *cache.Snapshot, loc protocol.Location) (changes []protocol.DocumentChange, _ error) {
+	pkg, pgf, err := NarrowestPackageForFile(ctx, snapshot, loc.URI)
+	if err != nil {
+		return nil, err
+	}
+
+	testBase := strings.TrimSuffix(filepath.Base(loc.URI.Path()), ".go") + "_test.go"
+	goTestFileURI := protocol.URIFromPath(filepath.Join(loc.URI.Dir().Path(), testBase))
+
+	testFH, err := snapshot.ReadFile(ctx, goTestFileURI)
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO(hxjiang): use a fresh name if the same test function name already
+	// exist.
+
+	var (
+		// edits contains all the text edits to be applied to the test file.
+		edits []protocol.TextEdit
+		// header is the buffer containing the text edit to the beginning of the file.
+		header bytes.Buffer
+	)
+
+	testPgf, err := snapshot.ParseGo(ctx, testFH, parsego.Header)
+	if err != nil {
+		if !errors.Is(err, os.ErrNotExist) {
+			return nil, err
+		}
+
+		changes = append(changes, protocol.DocumentChangeCreate(goTestFileURI))
+
+		// If this test file was created by the gopls, add a copyright header based
+		// on the originating file.
+		// Search for something that looks like a copyright header, to replicate
+		// in the new file.
+		// TODO(hxjiang): should we refine this heuristic, for example by checking for
+		// the word 'copyright'?
+		if groups := pgf.File.Comments; len(groups) > 0 {
+			// Copyright should appear before package decl and must be the first
+			// comment group.
+			// Avoid copying any other comment like package doc or directive comment.
+			if c := groups[0]; c.Pos() < pgf.File.Package && c != pgf.File.Doc &&
+				!isDirective(groups[0].List[0].Text) {
+				start, end, err := pgf.NodeOffsets(c)
+				if err != nil {
+					return nil, err
+				}
+				header.Write(pgf.Src[start:end])
+				// One empty line between copyright header and package decl.
+				header.WriteString("\n\n")
+			}
+		}
+	}
+
+	// If the test file does not have package decl, use the originating file to
+	// determine a package decl for the new file. Prefer xtest package.s
+	if testPgf == nil || testPgf.File == nil || testPgf.File.Package == token.NoPos {
+		// One empty line between package decl and rest of the file.
+		fmt.Fprintf(&header, "package %s_test\n\n", pkg.Types().Name())
+	}
+
+	// Write the copyright and package decl to the beginning of the file.
+	if text := header.String(); len(text) != 0 {
+		edits = append(edits, protocol.TextEdit{
+			Range:   protocol.Range{},
+			NewText: text,
+		})
+	}
+
+	// TODO(hxjiang): reject if the function/method is unexported.
+	// TODO(hxjiang): modify existing imports or add new imports.
+
+	// If the parse go file is missing, the fileEnd is the file start (zero value).
+	fileEnd := protocol.Range{}
+	if testPgf != nil {
+		fileEnd, err = testPgf.PosRange(testPgf.File.FileEnd, testPgf.File.FileEnd)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// test is the buffer containing the text edit to the test function.
+	var test bytes.Buffer
+	// TODO(hxjiang): replace test foo function with table-driven test.
+	test.WriteString("\nfunc TestFoo(*testing.T) {}")
+	edits = append(edits, protocol.TextEdit{
+		Range:   fileEnd,
+		NewText: test.String(),
+	})
+	return append(changes, protocol.DocumentChangeEdit(testFH, edits)), nil
+}

gopls/internal/golang/codeaction.go

@@ -227,6 +227,7 @@ type codeActionProducer struct {
 var codeActionProducers = [...]codeActionProducer{
 	{kind: protocol.QuickFix, fn: quickFix, needPkg: true},
 	{kind: protocol.SourceOrganizeImports, fn: sourceOrganizeImports},
+	{kind: settings.AddTest, fn: addTest, needPkg: true},
 	{kind: settings.GoAssembly, fn: goAssembly, needPkg: true},
 	{kind: settings.GoDoc, fn: goDoc, needPkg: true},
 	{kind: settings.GoFreeSymbols, fn: goFreeSymbols},
@@ -467,6 +468,41 @@ func refactorExtractToNewFile(ctx context.Context, req *codeActionsRequest) erro
 	return nil
 }
 
+// addTest produces "Add a test for FUNC" code actions.
+// See [server.commandHandler.AddTest] for command implementation.
+func addTest(ctx context.Context, req *codeActionsRequest) error {
+	// Reject if the feature is turned off.
+	if !req.snapshot.Options().AddTestSourceCodeAction {
+		return nil
+	}
+
+	// Reject test package.
+	if req.pkg.Metadata().ForTest != "" {
+		return nil
+	}
+
+	path, _ := astutil.PathEnclosingInterval(req.pgf.File, req.start, req.end)
+	if len(path) < 2 {
+		return nil
+	}
+
+	decl, ok := path[len(path)-2].(*ast.FuncDecl)
+	if !ok {
+		return nil
+	}
+
+	// Don't offer to create tests of "init" or "_".
+	if decl.Name.Name == "_" || decl.Name.Name == "init" {
+		return nil
+	}
+
+	cmd := command.NewAddTestCommand("Add a test for "+decl.Name.String(), req.loc)
+	req.addCommandAction(cmd, true)
+
+	// TODO(hxjiang): add code action for generate test for package/file.
+	return nil
+}
+
 // refactorRewriteRemoveUnusedParam produces "Remove unused parameter" code actions.
 // See [server.commandHandler.ChangeSignature] for command implementation.
 func refactorRewriteRemoveUnusedParam(ctx context.Context, req *codeActionsRequest) error {

gopls/internal/golang/extracttofile.go

@@ -80,7 +80,7 @@ func findImportEdits(file *ast.File, info *types.Info, start, end token.Pos) (ad
 }
 
 // ExtractToNewFile moves selected declarations into a new file.
-func ExtractToNewFile(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle, rng protocol.Range) (*protocol.WorkspaceEdit, error) {
+func ExtractToNewFile(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle, rng protocol.Range) ([]protocol.DocumentChange, error) {
 	errorPrefix := "ExtractToNewFile"
 
 	pkg, pgf, err := NarrowestPackageForFile(ctx, snapshot, fh.URI())
@@ -160,15 +160,15 @@ func ExtractToNewFile(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 		return nil, err
 	}
 
-	return protocol.NewWorkspaceEdit(
+	return []protocol.DocumentChange{
 		// edit the original file
 		protocol.DocumentChangeEdit(fh, append(importDeletes, protocol.TextEdit{Range: replaceRange, NewText: ""})),
 		// create a new file
 		protocol.DocumentChangeCreate(newFile.URI()),
 		// edit the created file
 		protocol.DocumentChangeEdit(newFile, []protocol.TextEdit{
 			{Range: protocol.Range{}, NewText: string(newFileContent)},
-		})), nil
+		})}, nil
 }
 
 // chooseNewFile chooses a new filename in dir, based on the name of the

gopls/internal/protocol/command/interface.go

@@ -224,6 +224,9 @@ type Interface interface {
 	// to avoid conflicts with other counters gopls collects.
 	AddTelemetryCounters(context.Context, AddTelemetryCountersArgs) error
 
+	// AddTest: add a test for the selected function
+	AddTest(context.Context, protocol.Location) (*protocol.WorkspaceEdit, error)
+
 	// MaybePromptForTelemetry: Prompt user to enable telemetry
 	//
 	// Checks for the right conditions, and then prompts the user

gopls/internal/server/command.go

@@ -275,6 +275,24 @@ func (*commandHandler) AddTelemetryCounters(_ context.Context, args command.AddT
 	return nil
 }
 
+func (c *commandHandler) AddTest(ctx context.Context, loc protocol.Location) (*protocol.WorkspaceEdit, error) {
+	var result *protocol.WorkspaceEdit
+	err := c.run(ctx, commandConfig{
+		forURI: loc.URI,
+	}, func(ctx context.Context, deps commandDeps) error {
+		if deps.snapshot.FileKind(deps.fh) != file.Go {
+			return fmt.Errorf("can't add test for non-Go file")
+		}
+		docedits, err := golang.AddTestForFunc(ctx, deps.snapshot, loc)
+		if err != nil {
+			return err
+		}
+		return applyChanges(ctx, c.s.client, docedits)
+	})
+	// TODO(hxjiang): move the cursor to the new test once edits applied.
+	return result, err
+}
+
 // commandConfig configures common command set-up and execution.
 type commandConfig struct {
 	requireSave bool                 // whether all files must be saved for the command to work
@@ -388,16 +406,7 @@ func (c *commandHandler) ApplyFix(ctx context.Context, args command.ApplyFixArgs
 			result = wsedit
 			return nil
 		}
-		resp, err := c.s.client.ApplyEdit(ctx, &protocol.ApplyWorkspaceEditParams{
-			Edit: *wsedit,
-		})
-		if err != nil {
-			return err
-		}
-		if !resp.Applied {
-			return errors.New(resp.FailureReason)
-		}
-		return nil
+		return applyChanges(ctx, c.s.client, changes)
 	})
 	return result, err
 }
@@ -622,17 +631,7 @@ func (c *commandHandler) RemoveDependency(ctx context.Context, args command.Remo
 		if err != nil {
 			return err
 		}
-		response, err := c.s.client.ApplyEdit(ctx, &protocol.ApplyWorkspaceEditParams{
-			Edit: *protocol.NewWorkspaceEdit(
-				protocol.DocumentChangeEdit(deps.fh, edits)),
-		})
-		if err != nil {
-			return err
-		}
-		if !response.Applied {
-			return fmt.Errorf("edits not applied because of %s", response.FailureReason)
-		}
-		return nil
+		return applyChanges(ctx, c.s.client, []protocol.DocumentChange{protocol.DocumentChangeEdit(deps.fh, edits)})
 	})
 }
 
@@ -1107,17 +1106,7 @@ func (c *commandHandler) AddImport(ctx context.Context, args command.AddImportAr
 		if err != nil {
 			return fmt.Errorf("could not add import: %v", err)
 		}
-		r, err := c.s.client.ApplyEdit(ctx, &protocol.ApplyWorkspaceEditParams{
-			Edit: *protocol.NewWorkspaceEdit(
-				protocol.DocumentChangeEdit(deps.fh, edits)),
-		})
-		if err != nil {
-			return fmt.Errorf("could not apply import edits: %v", err)
-		}
-		if !r.Applied {
-			return fmt.Errorf("failed to apply edits: %v", r.FailureReason)
-		}
-		return nil
+		return applyChanges(ctx, c.s.client, []protocol.DocumentChange{protocol.DocumentChangeEdit(deps.fh, edits)})
 	})
 }
 
@@ -1126,18 +1115,11 @@ func (c *commandHandler) ExtractToNewFile(ctx context.Context, args protocol.Loc
 		progress: "Extract to a new file",
 		forURI:   args.URI,
 	}, func(ctx context.Context, deps commandDeps) error {
-		edit, err := golang.ExtractToNewFile(ctx, deps.snapshot, deps.fh, args.Range)
+		changes, err := golang.ExtractToNewFile(ctx, deps.snapshot, deps.fh, args.Range)
 		if err != nil {
 			return err
 		}
-		resp, err := c.s.client.ApplyEdit(ctx, &protocol.ApplyWorkspaceEditParams{Edit: *edit})
-		if err != nil {
-			return fmt.Errorf("could not apply edits: %v", err)
-		}
-		if !resp.Applied {
-			return fmt.Errorf("edits not applied: %s", resp.FailureReason)
-		}
-		return nil
+		return applyChanges(ctx, c.s.client, changes)
 	})
 }
 
@@ -1543,13 +1525,7 @@ func (c *commandHandler) ChangeSignature(ctx context.Context, args command.Chang
 			result = wsedit
 			return nil
 		}
-		r, err := c.s.client.ApplyEdit(ctx, &protocol.ApplyWorkspaceEditParams{
-			Edit: *wsedit,
-		})
-		if !r.Applied {
-			return fmt.Errorf("failed to apply edits: %v", r.FailureReason)
-		}
-		return nil
+		return applyChanges(ctx, c.s.client, docedits)
 	})
 	return result, err
 }

gopls/internal/settings/codeactionkind.go

@@ -79,6 +79,7 @@ const (
 	GoDoc         protocol.CodeActionKind = "source.doc"
 	GoFreeSymbols protocol.CodeActionKind = "source.freesymbols"
 	GoTest        protocol.CodeActionKind = "source.test"
+	AddTest       protocol.CodeActionKind = "source.addTest"
 
 	// gopls
 	GoplsDocFeatures protocol.CodeActionKind = "gopls.doc.features"

gopls/internal/settings/default.go

@@ -136,6 +136,7 @@ func DefaultOptions(overrides ...func(*Options)) *Options {
 				LinkifyShowMessage:          false,
 				IncludeReplaceInWorkspace:   false,
 				ZeroConfig:                  true,
+				AddTestSourceCodeAction:     false,
 			},
 		}
 	})