internal/lsp: add list_known_packages and add_import commands

This CL adds two new commands that let a client request a list of importable packages relative to a Go file and then select which import a programmer would like to add to said file.

Updates golang/go#43351

Change-Id: If12518874a92ed4167bdd711a92e03ee21c7b949
Reviewed-on: https://go-review.googlesource.com/c/tools/+/281412
Trust: Rebecca Stambler <rstambler@golang.org>
Trust: Robert Findley <rfindley@google.com>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

gopls/doc/commands.md

@@ -21,16 +21,21 @@ Args:
 }
 ```
 
-### ****
+### **asks the server to add an import path to a given Go file.**
 Identifier: `gopls.add_import`
 
-
+The method will call applyEdit on the client so that clients don't have
+to apply the edit themselves.
 
 Args:
 
 ```
 {
+	// ImportPath is the target import path that should
+	// be added to the URI file
 	"ImportPath": string,
+	// URI is the file that the ImportPath should be
+	// added to
 	"URI": string,
 }
 ```
@@ -136,10 +141,10 @@ Args:
 }
 ```
 
-### ****
+### **retrieves a list of packages**
 Identifier: `gopls.list_known_packages`
 
-
+that are importable from the given URI.
 
 Args:
 

gopls/internal/regtest/misc/add_import_test.go

@@ -0,0 +1,59 @@
+// Copyright 2021 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 misc
+
+import (
+	"testing"
+
+	"golang.org/x/tools/internal/lsp/command"
+	"golang.org/x/tools/internal/lsp/protocol"
+	. "golang.org/x/tools/internal/lsp/regtest"
+	"golang.org/x/tools/internal/lsp/tests"
+)
+
+func TestAddImport(t *testing.T) {
+	const before = `package main
+
+import "fmt"
+
+func main() {
+	fmt.Println("hello world")
+}
+`
+
+	const want = `package main
+
+import (
+	"bytes"
+	"fmt"
+)
+
+func main() {
+	fmt.Println("hello world")
+}
+`
+
+	Run(t, "", func(t *testing.T, env *Env) {
+		env.CreateBuffer("main.go", before)
+		cmd, err := command.NewAddImportCommand("Add Import", command.AddImportArgs{
+			URI:        protocol.URIFromSpanURI(env.Sandbox.Workdir.URI("main.go").SpanURI()),
+			ImportPath: "bytes",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		_, err = env.Editor.ExecuteCommand(env.Ctx, &protocol.ExecuteCommandParams{
+			Command:   "gopls.add_import",
+			Arguments: cmd.Arguments,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		got := env.Editor.BufferText("main.go")
+		if got != want {
+			t.Fatalf("gopls.add_import failed\n%s", tests.Diff(t, want, got))
+		}
+	})
+}

internal/lsp/cache/check.go

@@ -503,7 +503,7 @@ func doTypeCheck(ctx context.Context, snapshot *snapshot, m *metadata, mode sour
 			if dep == nil {
 				return nil, snapshot.missingPkgError(pkgPath)
 			}
-			if !isValidImport(m.pkgPath, dep.m.pkgPath) {
+			if !source.IsValidImport(string(m.pkgPath), string(dep.m.pkgPath)) {
 				return nil, errors.Errorf("invalid use of internal package %s", pkgPath)
 			}
 			depPkg, err := dep.check(ctx, snapshot)
@@ -834,17 +834,6 @@ func resolveImportPath(importPath string, pkg *pkg, deps map[packagePath]*packag
 	return nil
 }
 
-func isValidImport(pkgPath, importPkgPath packagePath) bool {
-	i := strings.LastIndex(string(importPkgPath), "/internal/")
-	if i == -1 {
-		return true
-	}
-	if isCommandLineArguments(string(pkgPath)) {
-		return true
-	}
-	return strings.HasPrefix(string(pkgPath), string(importPkgPath[:i]))
-}
-
 // An importFunc is an implementation of the single-method
 // types.Importer interface based on a function value.
 type importerFunc func(path string) (*types.Package, error)

internal/lsp/cache/load.go

@@ -55,7 +55,7 @@ func (s *snapshot) load(ctx context.Context, allowNetwork bool, scopes ...interf
 	for _, scope := range scopes {
 		switch scope := scope.(type) {
 		case packagePath:
-			if isCommandLineArguments(string(scope)) {
+			if source.IsCommandLineArguments(string(scope)) {
 				panic("attempted to load command-line-arguments")
 			}
 			// The only time we pass package paths is when we're doing a

internal/lsp/cache/snapshot.go

@@ -995,7 +995,7 @@ func (s *snapshot) addID(uri span.URI, id packageID) {
 		}
 		// If the package previously only had a command-line-arguments ID,
 		// we should just replace it.
-		if isCommandLineArguments(string(existingID)) {
+		if source.IsCommandLineArguments(string(existingID)) {
 			s.ids[uri][i] = id
 			// Delete command-line-arguments if it was a workspace package.
 			delete(s.workspacePackages, existingID)
@@ -1012,14 +1012,6 @@ func (s *snapshot) addID(uri span.URI, id packageID) {
 	s.ids[uri] = append(newIDs, id)
 }
 
-// isCommandLineArguments reports whether a given value denotes
-// "command-line-arguments" package, which is a package with an unknown ID
-// created by the go command. It can have a test variant, which is why callers
-// should not check that a value equals "command-line-arguments" directly.
-func isCommandLineArguments(s string) bool {
-	return strings.Contains(s, "command-line-arguments")
-}
-
 func (s *snapshot) isWorkspacePackage(id packageID) bool {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -1169,7 +1161,7 @@ func shouldShowAdHocPackagesWarning(snapshot source.Snapshot, pkgs []source.Pack
 
 func containsCommandLineArguments(pkgs []source.Package) bool {
 	for _, pkg := range pkgs {
-		if isCommandLineArguments(pkg.ID()) {
+		if source.IsCommandLineArguments(pkg.ID()) {
 			return true
 		}
 	}
@@ -1237,7 +1229,7 @@ func (s *snapshot) reloadWorkspace(ctx context.Context) error {
 		missingMetadata = true
 
 		// Don't try to reload "command-line-arguments" directly.
-		if isCommandLineArguments(string(pkgPath)) {
+		if source.IsCommandLineArguments(string(pkgPath)) {
 			continue
 		}
 		pkgPathSet[pkgPath] = struct{}{}
@@ -1646,7 +1638,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		// go command when the user is outside of GOPATH and outside of a
 		// module. Do not cache them as workspace packages for longer than
 		// necessary.
-		if isCommandLineArguments(string(id)) {
+		if source.IsCommandLineArguments(string(id)) {
 			if invalidateMetadata, ok := idsToInvalidate[id]; invalidateMetadata && ok {
 				continue
 			}

internal/lsp/cmd/test/cmdtest.go

@@ -100,6 +100,10 @@ func (r *runner) FunctionExtraction(t *testing.T, start span.Span, end span.Span
 	//TODO: function extraction not supported on command line
 }
 
+func (r *runner) AddImport(t *testing.T, uri span.URI, expectedImport string) {
+	//TODO: import addition not supported on command line
+}
+
 func (r *runner) runGoplsCmd(t testing.TB, args ...string) (string, string) {
 	rStdout, wStdout, err := os.Pipe()
 	if err != nil {

internal/lsp/command.go

@@ -54,7 +54,7 @@ type commandHandler struct {
 
 // commandConfig configures common command set-up and execution.
 type commandConfig struct {
-	async       bool                 // whether to run the command asynchronously. Async commands cannot return results.
+	async       bool                 // whether to run the command asynchronously. Async commands can only return errors.
 	requireSave bool                 // whether all files must be saved for the command to work
 	progress    string               // title to use for progress reporting. If empty, no progress will be reported.
 	forURI      protocol.DocumentURI // URI to resolve to a snapshot. If unset, snapshot will be nil.
@@ -108,7 +108,16 @@ func (c *commandHandler) run(ctx context.Context, cfg commandConfig, run command
 		return err
 	}
 	if cfg.async {
-		go runcmd()
+		go func() {
+			if err := runcmd(); err != nil {
+				if showMessageErr := c.s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
+					Type:    protocol.Error,
+					Message: err.Error(),
+				}); showMessageErr != nil {
+					event.Error(ctx, fmt.Sprintf("failed to show message: %q", err.Error()), showMessageErr)
+				}
+			}
+		}()
 		return nil
 	}
 	return runcmd()
@@ -335,15 +344,8 @@ func (c *commandHandler) RunTests(ctx context.Context, args command.RunTestsArgs
 		forURI:      args.URI,
 	}, func(ctx context.Context, deps commandDeps) error {
 		if err := c.runTests(ctx, deps.snapshot, deps.work, args.URI, args.Tests, args.Benchmarks); err != nil {
-			if err := c.s.client.ShowMessage(ctx, &protocol.ShowMessageParams{
-				Type:    protocol.Error,
-				Message: fmt.Sprintf("Running tests failed: %v", err),
-			}); err != nil {
-				event.Error(ctx, "running tests: failed to show message", err)
-			}
+			return errors.Errorf("running tests failed: %w", err)
 		}
-		// Since we're running asynchronously, any error returned here would be
-		// ignored.
 		return nil
 	})
 }
@@ -664,26 +666,33 @@ func (c *commandHandler) GenerateGoplsMod(ctx context.Context, args command.URIA
 func (c *commandHandler) ListKnownPackages(ctx context.Context, args command.URIArg) (command.ListKnownPackagesResult, error) {
 	var result command.ListKnownPackagesResult
 	err := c.run(ctx, commandConfig{
-		progress: "Listing packages", // optional, causes a progress report during command execution
-		forURI:   args.URI,           // optional, populates deps.snapshot and deps.fh
+		progress: "Listing packages",
+		forURI:   args.URI,
 	}, func(ctx context.Context, deps commandDeps) error {
-		// Marwan: add implementation here. deps.snapshot and deps.fh are available for use.
-		result.Packages = []string{}
-		return nil
+		var err error
+		result.Packages, err = source.KnownPackages(ctx, deps.snapshot, deps.fh)
+		return err
 	})
 	return result, err
 }
-
-func (c *commandHandler) AddImport(ctx context.Context, args command.AddImportArgs) (command.AddImportResult, error) {
-	var result command.AddImportResult
-	err := c.run(ctx, commandConfig{
+func (c *commandHandler) AddImport(ctx context.Context, args command.AddImportArgs) error {
+	return c.run(ctx, commandConfig{
 		progress: "Adding import",
 		forURI:   args.URI,
 	}, func(ctx context.Context, deps commandDeps) error {
-		result.Edits = nil
+		edits, err := source.AddImport(ctx, deps.snapshot, deps.fh, args.ImportPath)
+		if err != nil {
+			return fmt.Errorf("could not add import: %v", err)
+		}
+		if _, err := c.s.client.ApplyEdit(ctx, &protocol.ApplyWorkspaceEditParams{
+			Edit: protocol.WorkspaceEdit{
+				DocumentChanges: documentChanges(deps.fh, edits),
+			},
+		}); err != nil {
+			return fmt.Errorf("could not apply import edits: %v", err)
+		}
 		return nil
 	})
-	return result, err
 }
 
 func (c *commandHandler) WorkspaceMetadata(ctx context.Context) (command.WorkspaceMetadataResult, error) {

internal/lsp/command/command_gen.go

@@ -77,7 +77,7 @@ func Dispatch(ctx context.Context, params *protocol.ExecuteCommandParams, s Inte
 		if err := UnmarshalArgs(params.Arguments, &a0); err != nil {
 			return nil, err
 		}
-		return s.AddImport(ctx, a0)
+		return nil, s.AddImport(ctx, a0)
 	case "gopls.apply_fix":
 		var a0 ApplyFixArgs
 		if err := UnmarshalArgs(params.Arguments, &a0); err != nil {

internal/lsp/command/interface.go

@@ -115,9 +115,14 @@ type Interface interface {
 	// (Re)generate the gopls.mod file for a workspace.
 	GenerateGoplsMod(context.Context, URIArg) error
 
+	// ListKnownPackages: retrieves a list of packages
+	// that are importable from the given URI.
 	ListKnownPackages(context.Context, URIArg) (ListKnownPackagesResult, error)
 
-	AddImport(context.Context, AddImportArgs) (AddImportResult, error)
+	// AddImport: asks the server to add an import path to a given Go file.
+	// The method will call applyEdit on the client so that clients don't have
+	// to apply the edit themselves.
+	AddImport(context.Context, AddImportArgs) error
 
 	WorkspaceMetadata(context.Context) (WorkspaceMetadataResult, error)
 
@@ -196,18 +201,21 @@ type GoGetPackageArgs struct {
 	AddRequire bool
 }
 
-// TODO (Marwan): document :)
-
 type AddImportArgs struct {
+	// ImportPath is the target import path that should
+	// be added to the URI file
 	ImportPath string
-	URI        protocol.DocumentURI
-}
-
-type AddImportResult struct {
-	Edits []protocol.TextDocumentEdit
+	// URI is the file that the ImportPath should be
+	// added to
+	URI protocol.DocumentURI
 }
 
 type ListKnownPackagesResult struct {
+	// Packages is a list of packages relative
+	// to the URIArg passed by the command request.
+	// In other words, it omits paths that are already
+	// imported or cannot be imported due to compiler
+	// restrictions.
 	Packages []string
 }
 

internal/lsp/lsp_test.go

@@ -16,6 +16,7 @@ import (
 	"testing"
 
 	"golang.org/x/tools/internal/lsp/cache"
+	"golang.org/x/tools/internal/lsp/command"
 	"golang.org/x/tools/internal/lsp/diff"
 	"golang.org/x/tools/internal/lsp/diff/myers"
 	"golang.org/x/tools/internal/lsp/protocol"
@@ -111,6 +112,10 @@ func (c testClient) PublishDiagnostics(context.Context, *protocol.PublishDiagnos
 	return nil
 }
 
+func (c testClient) ShowMessage(context.Context, *protocol.ShowMessageParams) error {
+	return nil
+}
+
 func (c testClient) ApplyEdit(ctx context.Context, params *protocol.ApplyWorkspaceEditParams) (*protocol.ApplyWorkspaceEditResponse, error) {
 	res, err := applyTextDocumentEdits(c.runner, params.Edit.DocumentChanges)
 	if err != nil {
@@ -1098,6 +1103,57 @@ func (r *runner) Link(t *testing.T, uri span.URI, wantLinks []tests.Link) {
 	}
 }
 
+func (r *runner) AddImport(t *testing.T, uri span.URI, expectedImport string) {
+	cmd, err := command.NewListKnownPackagesCommand("List Known Packages", command.URIArg{
+		URI: protocol.URIFromSpanURI(uri),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp, err := r.server.executeCommand(r.ctx, &protocol.ExecuteCommandParams{
+		Command:   cmd.Command,
+		Arguments: cmd.Arguments,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	res := resp.(command.ListKnownPackagesResult)
+	var hasPkg bool
+	for _, p := range res.Packages {
+		if p == expectedImport {
+			hasPkg = true
+			break
+		}
+	}
+	if !hasPkg {
+		t.Fatalf("%s: got %v packages\nwant contains %q", command.ListKnownPackages, res.Packages, expectedImport)
+	}
+	cmd, err = command.NewAddImportCommand("Add Imports", command.AddImportArgs{
+		URI:        protocol.URIFromSpanURI(uri),
+		ImportPath: expectedImport,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = r.server.executeCommand(r.ctx, &protocol.ExecuteCommandParams{
+		Command:   cmd.Command,
+		Arguments: cmd.Arguments,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := (<-r.editRecv)[uri]
+	want := r.data.Golden("addimport", uri.Filename(), func() ([]byte, error) {
+		return []byte(got), nil
+	})
+	if want == nil {
+		t.Fatalf("golden file %q not found", uri.Filename())
+	}
+	if diff := tests.Diff(t, got, string(want)); diff != "" {
+		t.Errorf("%s mismatch\n%s", command.AddImport, diff)
+	}
+}
+
 func TestBytesOffset(t *testing.T) {
 	tests := []struct {
 		text string

internal/lsp/source/add_import.go

@@ -0,0 +1,26 @@
+// 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 source
+
+import (
+	"context"
+
+	"golang.org/x/tools/internal/imports"
+	"golang.org/x/tools/internal/lsp/protocol"
+)
+
+// AddImport adds a single import statement to the given file
+func AddImport(ctx context.Context, snapshot Snapshot, fh VersionedFileHandle, importPath string) ([]protocol.TextEdit, error) {
+	_, pgf, err := GetParsedFile(ctx, snapshot, fh, NarrowestPackage)
+	if err != nil {
+		return nil, err
+	}
+	return ComputeOneImportFixEdits(snapshot, pgf, &imports.ImportFix{
+		StmtInfo: imports.ImportInfo{
+			ImportPath: importPath,
+		},
+		FixType: imports.AddImport,
+	})
+}

internal/lsp/source/api_json.go

@@ -751,9 +751,9 @@ var GeneratedAPIJSON = &APIJSON{
 		},
 		{
 			Command: "gopls.add_import",
-			Title:   "",
-			Doc:     "",
-			ArgDoc:  "{\n\t\"ImportPath\": string,\n\t\"URI\": string,\n}",
+			Title:   "asks the server to add an import path to a given Go file.",
+			Doc:     "The method will call applyEdit on the client so that clients don't have\nto apply the edit themselves.",
+			ArgDoc:  "{\n\t// ImportPath is the target import path that should\n\t// be added to the URI file\n\t\"ImportPath\": string,\n\t// URI is the file that the ImportPath should be\n\t// added to\n\t\"URI\": string,\n}",
 		},
 		{
 			Command: "gopls.apply_fix",
@@ -793,8 +793,8 @@ var GeneratedAPIJSON = &APIJSON{
 		},
 		{
 			Command: "gopls.list_known_packages",
-			Title:   "",
-			Doc:     "",
+			Title:   "retrieves a list of packages",
+			Doc:     "that are importable from the given URI.",
 			ArgDoc:  "{\n\t// The file URI.\n\t\"URI\": string,\n}",
 		},
 		{

internal/lsp/source/known_packages.go

@@ -0,0 +1,118 @@
+// 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 source
+
+import (
+	"context"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/imports"
+	errors "golang.org/x/xerrors"
+)
+
+// KnownPackages returns a list of all known packages
+// in the package graph that could potentially be imported
+// by the given file.
+func KnownPackages(ctx context.Context, snapshot Snapshot, fh VersionedFileHandle) ([]string, error) {
+	pkg, pgf, err := GetParsedFile(ctx, snapshot, fh, NarrowestPackage)
+	if err != nil {
+		return nil, errors.Errorf("GetParsedFile: %w", err)
+	}
+	alreadyImported := map[string]struct{}{}
+	for _, imp := range pgf.File.Imports {
+		alreadyImported[imp.Path.Value] = struct{}{}
+	}
+	pkgs, err := snapshot.CachedImportPaths(ctx)
+	if err != nil {
+		return nil, err
+	}
+	var (
+		seen  = make(map[string]struct{})
+		paths []string
+	)
+	for path, knownPkg := range pkgs {
+		gofiles := knownPkg.CompiledGoFiles()
+		if len(gofiles) == 0 || gofiles[0].File.Name == nil {
+			continue
+		}
+		pkgName := gofiles[0].File.Name.Name
+		// package main cannot be imported
+		if pkgName == "main" {
+			continue
+		}
+		// test packages cannot be imported
+		if knownPkg.ForTest() != "" {
+			continue
+		}
+		// no need to import what the file already imports
+		if _, ok := alreadyImported[path]; ok {
+			continue
+		}
+		// snapshot.KnownPackages could have multiple versions of a pkg
+		if _, ok := seen[path]; ok {
+			continue
+		}
+		seen[path] = struct{}{}
+		// make sure internal packages are importable by the file
+		if !IsValidImport(pkg.PkgPath(), path) {
+			continue
+		}
+		// naive check on cyclical imports
+		if isDirectlyCyclical(pkg, knownPkg) {
+			continue
+		}
+		paths = append(paths, path)
+		seen[path] = struct{}{}
+	}
+	err = snapshot.RunProcessEnvFunc(ctx, func(o *imports.Options) error {
+		var mu sync.Mutex
+		ctx, cancel := context.WithTimeout(ctx, time.Millisecond*80)
+		defer cancel()
+		return imports.GetAllCandidates(ctx, func(ifix imports.ImportFix) {
+			mu.Lock()
+			defer mu.Unlock()
+			if _, ok := seen[ifix.StmtInfo.ImportPath]; ok {
+				return
+			}
+			paths = append(paths, ifix.StmtInfo.ImportPath)
+		}, "", pgf.URI.Filename(), pkg.GetTypes().Name(), o.Env)
+	})
+	if err != nil {
+		// if an error occurred, we stil have a decent list we can
+		// show to the user through snapshot.CachedImportPaths
+		event.Error(ctx, "imports.GetAllCandidates", err)
+	}
+	sort.Slice(paths, func(i, j int) bool {
+		importI, importJ := paths[i], paths[j]
+		iHasDot := strings.Contains(importI, ".")
+		jHasDot := strings.Contains(importJ, ".")
+		if iHasDot && !jHasDot {
+			return false
+		}
+		if jHasDot && !iHasDot {
+			return true
+		}
+		return importI < importJ
+	})
+	return paths, nil
+}
+
+// isDirectlyCyclical checks if imported directly imports pkg.
+// It does not (yet) offer a full cyclical check because showing a user
+// a list of importable packages already generates a very large list
+// and having a few false positives in there could be worth the
+// performance snappiness.
+func isDirectlyCyclical(pkg, imported Package) bool {
+	for _, imp := range imported.Imports() {
+		if imp.PkgPath() == pkg.PkgPath() {
+			return true
+		}
+	}
+	return false
+}

internal/lsp/source/source_test.go

@@ -936,6 +936,7 @@ func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string,
 }
 func (r *runner) FunctionExtraction(t *testing.T, start span.Span, end span.Span) {}
 func (r *runner) CodeLens(t *testing.T, uri span.URI, want []protocol.CodeLens)   {}
+func (r *runner) AddImport(t *testing.T, uri span.URI, expectedImport string)     {}
 
 func spanToRange(data *tests.Data, spn span.Span) (*protocol.ColumnMapper, protocol.Range, error) {
 	m, err := data.Mapper(spn.URI())

internal/lsp/source/util.go

@@ -507,3 +507,24 @@ func inDirLex(dir, path string) bool {
 		return false
 	}
 }
+
+// IsValidImport returns whether importPkgPath is importable
+// by pkgPath
+func IsValidImport(pkgPath, importPkgPath string) bool {
+	i := strings.LastIndex(string(importPkgPath), "/internal/")
+	if i == -1 {
+		return true
+	}
+	if IsCommandLineArguments(string(pkgPath)) {
+		return true
+	}
+	return strings.HasPrefix(string(pkgPath), string(importPkgPath[:i]))
+}
+
+// IsCommandLineArguments reports whether a given value denotes
+// "command-line-arguments" package, which is a package with an unknown ID
+// created by the go command. It can have a test variant, which is why callers
+// should not check that a value equals "command-line-arguments" directly.
+func IsCommandLineArguments(s string) bool {
+	return strings.Contains(s, "command-line-arguments")
+}

internal/lsp/testdata/addimport/addimport.go.golden

@@ -0,0 +1,7 @@
+-- addimport --
+package addimport //@addimport("", "bytes")
+
+import "bytes"
+
+func main() {}
+

internal/lsp/testdata/addimport/addimport.go.in

@@ -0,0 +1,3 @@
+package addimport //@addimport("", "bytes")
+
+func main() {}

internal/lsp/tests/tests.go

@@ -74,6 +74,7 @@ type SymbolInformation map[span.Span]protocol.SymbolInformation
 type WorkspaceSymbols map[WorkspaceSymbolsTestType]map[span.URI][]string
 type Signatures map[span.Span]*protocol.SignatureHelp
 type Links map[span.URI][]Link
+type AddImport map[span.URI]string
 
 type Data struct {
 	Config                   packages.Config
@@ -107,6 +108,7 @@ type Data struct {
 	WorkspaceSymbols         WorkspaceSymbols
 	Signatures               Signatures
 	Links                    Links
+	AddImport                AddImport
 
 	t         testing.TB
 	fragments map[string]string
@@ -147,6 +149,7 @@ type Tests interface {
 	WorkspaceSymbols(*testing.T, span.URI, string, WorkspaceSymbolsTestType)
 	SignatureHelp(*testing.T, span.Span, *protocol.SignatureHelp)
 	Link(*testing.T, span.URI, []Link)
+	AddImport(*testing.T, span.URI, string)
 }
 
 type Definition struct {
@@ -293,6 +296,7 @@ func load(t testing.TB, mode string, dir string) *Data {
 		WorkspaceSymbols:         make(WorkspaceSymbols),
 		Signatures:               make(Signatures),
 		Links:                    make(Links),
+		AddImport:                make(AddImport),
 
 		t:         t,
 		dir:       dir,
@@ -447,6 +451,7 @@ func load(t testing.TB, mode string, dir string) *Data {
 		"extractfunc":     datum.collectFunctionExtractions,
 		"incomingcalls":   datum.collectIncomingCalls,
 		"outgoingcalls":   datum.collectOutgoingCalls,
+		"addimport":       datum.collectAddImports,
 	}); err != nil {
 		t.Fatal(err)
 	}
@@ -786,6 +791,15 @@ func Run(t *testing.T, tests Tests, data *Data) {
 		}
 	})
 
+	t.Run("AddImport", func(t *testing.T) {
+		t.Helper()
+		for uri, exp := range data.AddImport {
+			t.Run(uriName(uri), func(t *testing.T) {
+				tests.AddImport(t, uri, exp)
+			})
+		}
+	})
+
 	if *UpdateGolden {
 		for _, golden := range data.golden {
 			if !golden.Modified {
@@ -1077,6 +1091,10 @@ func (data *Data) collectImports(spn span.Span) {
 	data.Imports = append(data.Imports, spn)
 }
 
+func (data *Data) collectAddImports(spn span.Span, imp string) {
+	data.AddImport[spn.URI()] = imp
+}
+
 func (data *Data) collectSemanticTokens(spn span.Span) {
 	data.SemanticTokens = append(data.SemanticTokens, spn)
 }