internal/lsp/regtest: add support for custom test proxy data

Certain regtests require referencing external data. To support this, add
the ability to use a file-based proxy populated with testdata.

To expose this configuration, augment the regtest runner with variadic
options. Also use this to replace the Runner.RunInMode function.

Add a simple regtest that uses this functionality.

Updates golang/go#36879

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

7 files changed, 199 insertions, 23 deletions

diff --git a/internal/lsp/fake/editor_test.go b/internal/lsp/fake/editor_test.go
index 96575530e..cc08d8dbf 100644
--- a/internal/lsp/fake/editor_test.go
+++ b/internal/lsp/fake/editor_test.go
@@ -48,7 +48,7 @@ func main() {
 `
 
 func TestClientEditing(t *testing.T) {
-	ws, err := NewWorkspace("TestClientEditing", []byte(exampleProgram))
+	ws, err := NewWorkspace("TestClientEditing", exampleProgram, "")
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/lsp/fake/workspace.go b/internal/lsp/fake/workspace.go
index 7e20ff89e..f3c4a1b05 100644
--- a/internal/lsp/fake/workspace.go
+++ b/internal/lsp/fake/workspace.go
@@ -15,6 +15,7 @@ import (
 
 	"golang.org/x/tools/internal/gocommand"
 	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/proxydir"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/txtar"
 )
@@ -29,9 +30,10 @@ type FileEvent struct {
 // The Workspace type represents a temporary workspace to use for editing Go
 // files in tests.
 type Workspace struct {
-	name    string
-	gopath  string
-	workdir string
+	name     string
+	gopath   string
+	workdir  string
+	proxydir string
 
 	watcherMu sync.Mutex
 	watchers  []func(context.Context, []FileEvent)
@@ -40,7 +42,7 @@ type Workspace struct {
 // NewWorkspace creates a named workspace populated by the txtar-encoded
 // content given by txt. It creates temporary directories for the workspace
 // content and for GOPATH.
-func NewWorkspace(name string, txt []byte) (_ *Workspace, err error) {
+func NewWorkspace(name string, srctxt string, proxytxt string) (_ *Workspace, err error) {
 	w := &Workspace{name: name}
 	defer func() {
 		// Clean up if we fail at any point in this constructor.
@@ -58,15 +60,73 @@ func NewWorkspace(name string, txt []byte) (_ *Workspace, err error) {
 		return nil, fmt.Errorf("creating temporary gopath: %v", err)
 	}
 	w.gopath = gopath
-	archive := txtar.Parse(txt)
-	for _, f := range archive.Files {
-		if err := w.writeFileData(f.Name, string(f.Data)); err != nil {
-			return nil, err
+	files := unpackTxt(srctxt)
+	for name, data := range files {
+		if err := w.writeFileData(name, string(data)); err != nil {
+			return nil, fmt.Errorf("writing to workdir: %v", err)
 		}
 	}
+	pd, err := ioutil.TempDir("", fmt.Sprintf("goplstest-proxy-%s-", name))
+	if err != nil {
+		return nil, fmt.Errorf("creating temporary proxy dir: %v", err)
+	}
+	w.proxydir = pd
+	if err := writeProxyDir(unpackTxt(proxytxt), w.proxydir); err != nil {
+		return nil, fmt.Errorf("writing proxy dir: %v", err)
+	}
 	return w, nil
 }
 
+func unpackTxt(txt string) map[string][]byte {
+	dataMap := make(map[string][]byte)
+	archive := txtar.Parse([]byte(txt))
+	for _, f := range archive.Files {
+		dataMap[f.Name] = f.Data
+	}
+	return dataMap
+}
+
+func writeProxyDir(files map[string][]byte, dir string) error {
+	type moduleVersion struct {
+		modulePath, version string
+	}
+	// Transform into the format expected by the proxydir package.
+	filesByModule := make(map[moduleVersion]map[string][]byte)
+	for name, data := range files {
+		modulePath, version, suffix := splitModuleVersionPath(name)
+		mv := moduleVersion{modulePath, version}
+		if _, ok := filesByModule[mv]; !ok {
+			filesByModule[mv] = make(map[string][]byte)
+		}
+		filesByModule[mv][suffix] = data
+	}
+	for mv, files := range filesByModule {
+		if err := proxydir.WriteModuleVersion(dir, mv.modulePath, mv.version, files); err != nil {
+			return fmt.Errorf("error writing %s@%s: %v", mv.modulePath, mv.version, err)
+		}
+	}
+	return nil
+}
+
+// splitModuleVersionPath extracts module information from files stored in the
+// directory structure modulePath@version/suffix.
+// For example:
+//  splitModuleVersionPath("mod.com@v1.2.3/package") = ("mod.com", "v1.2.3", "package")
+func splitModuleVersionPath(path string) (modulePath, version, suffix string) {
+	parts := strings.Split(path, "/")
+	var modulePathParts []string
+	for i, p := range parts {
+		if strings.Contains(p, "@") {
+			mv := strings.SplitN(p, "@", 2)
+			modulePathParts = append(modulePathParts, mv[0])
+			return strings.Join(modulePathParts, "/"), mv[1], strings.Join(parts[i+1:], "/")
+		}
+		modulePathParts = append(modulePathParts, p)
+	}
+	// Default behavior: this is just a module path.
+	return path, "", ""
+}
+
 // RootURI returns the root URI for this workspace.
 func (w *Workspace) RootURI() protocol.DocumentURI {
 	return toURI(w.workdir)
@@ -77,6 +137,11 @@ func (w *Workspace) GOPATH() string {
 	return w.gopath
 }
 
+// GOPROXY returns the value that GOPROXY should be set to for this workspace.
+func (w *Workspace) GOPROXY() string {
+	return proxydir.ToURL(w.proxydir)
+}
+
 // AddWatcher registers the given func to be called on any file change.
 func (w *Workspace) AddWatcher(watcher func(context.Context, []FileEvent)) {
 	w.watcherMu.Lock()
@@ -156,6 +221,7 @@ func (w *Workspace) RemoveFile(ctx context.Context, path string) error {
 func (w *Workspace) GoEnv() []string {
 	return []string{
 		"GOPATH=" + w.GOPATH(),
+		"GOPROXY=" + w.GOPROXY(),
 		"GO111MODULE=",
 		"GOSUMDB=off",
 	}
diff --git a/internal/lsp/fake/workspace_test.go b/internal/lsp/fake/workspace_test.go
index 31b46d2af..1ab0d658f 100644
--- a/internal/lsp/fake/workspace_test.go
+++ b/internal/lsp/fake/workspace_test.go
@@ -21,13 +21,13 @@ Hello World!
 func newWorkspace(t *testing.T) (*Workspace, <-chan []FileEvent, func()) {
 	t.Helper()
 
-	ws, err := NewWorkspace("default", []byte(data))
+	ws, err := NewWorkspace("default", data, "")
 	if err != nil {
 		t.Fatal(err)
 	}
 	cleanup := func() {
 		if err := ws.Close(); err != nil {
-			t.Fatal(err)
+			t.Errorf("closing workspace: %v", err)
 		}
 	}
 
@@ -90,3 +90,24 @@ func TestWorkspace_WriteFile(t *testing.T) {
 		}
 	}
 }
+
+func TestSplitModuleVersionPath(t *testing.T) {
+	tests := []struct {
+		path                                string
+		wantModule, wantVersion, wantSuffix string
+	}{
+		{"foo.com@v1.2.3/bar", "foo.com", "v1.2.3", "bar"},
+		{"foo.com/module@v1.2.3/bar", "foo.com/module", "v1.2.3", "bar"},
+		{"foo.com@v1.2.3", "foo.com", "v1.2.3", ""},
+		{"std@v1.14.0", "std", "v1.14.0", ""},
+		{"another/module/path", "another/module/path", "", ""},
+	}
+
+	for _, test := range tests {
+		module, version, suffix := splitModuleVersionPath(test.path)
+		if module != test.wantModule || version != test.wantVersion || suffix != test.wantSuffix {
+			t.Errorf("splitModuleVersionPath(%q) =\n\t(%q, %q, %q)\nwant\n\t(%q, %q, %q)",
+				test.path, module, version, suffix, test.wantModule, test.wantVersion, test.wantSuffix)
+		}
+	}
+}
diff --git a/internal/lsp/lsprpc/lsprpc_test.go b/internal/lsp/lsprpc/lsprpc_test.go
index 0f03e67c4..67db2145e 100644
--- a/internal/lsp/lsprpc/lsprpc_test.go
+++ b/internal/lsp/lsprpc/lsprpc_test.go
@@ -190,7 +190,7 @@ func TestDebugInfoLifecycle(t *testing.T) {
 	resetExitFuncs := OverrideExitFuncsForTest()
 	defer resetExitFuncs()
 
-	ws, err := fake.NewWorkspace("gopls-lsprpc-test", []byte(exampleProgram))
+	ws, err := fake.NewWorkspace("gopls-lsprpc-test", exampleProgram, "")
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/lsp/regtest/diagnostics_test.go b/internal/lsp/regtest/diagnostics_test.go
index 7555518ba..f78c32372 100644
--- a/internal/lsp/regtest/diagnostics_test.go
+++ b/internal/lsp/regtest/diagnostics_test.go
@@ -259,3 +259,48 @@ func TestPackageChange(t *testing.T) {
 		env.Await(EmptyDiagnostics("a.go"))
 	})
 }
+
+const testPackageWithRequire = `
+-- go.mod --
+module mod.com
+
+go 1.12
+
+require (
+	foo.test v1.2.3
+)
+-- print.go --
+package lib
+
+import (
+	"fmt"
+
+	"foo.test/bar"
+)
+
+func PrintAnswer() {
+	fmt.Printf("answer: %s", bar.Answer)
+}
+`
+
+const testPackageWithRequireProxy = `
+-- foo.test@v1.2.3/go.mod --
+module foo.test
+
+go 1.12
+-- foo.test@v1.2.3/bar/const.go --
+package bar
+
+const Answer = 42
+`
+
+func TestResolveDiagnosticWithDownload(t *testing.T) {
+	runner.Run(t, testPackageWithRequire, func(t *testing.T, env *Env) {
+		env.OpenFile("print.go")
+		env.W.RunGoCommand(env.Ctx, "mod", "download")
+		// Check that gopackages correctly loaded this dependency. We should get a
+		// diagnostic for the wrong formatting type.
+		// TODO: we should be able to easily also match the diagnostic message.
+		env.Await(env.DiagnosticAtRegexp("print.go", "fmt.Printf"))
+	}, WithProxy(testPackageWithRequireProxy))
+}
diff --git a/internal/lsp/regtest/env.go b/internal/lsp/regtest/env.go
index 827e6855d..2faf4256b 100644
--- a/internal/lsp/regtest/env.go
+++ b/internal/lsp/regtest/env.go
@@ -163,17 +163,61 @@ func (r *Runner) Close() error {
 	return nil
 }
 
+type runConfig struct {
+	modes    EnvMode
+	proxyTxt string
+	timeout  time.Duration
+}
+
+func (r *Runner) defaultConfig() *runConfig {
+	return &runConfig{
+		modes:   r.defaultModes,
+		timeout: r.timeout,
+	}
+}
+
+// A RunOption augments the behavior of the test runner.
+type RunOption interface {
+	set(*runConfig)
+}
+
+type optionSetter func(*runConfig)
+
+func (f optionSetter) set(opts *runConfig) {
+	f(opts)
+}
+
+// WithTimeout configures a custom timeout for this test run.
+func WithTimeout(d time.Duration) RunOption {
+	return optionSetter(func(opts *runConfig) {
+		opts.timeout = d
+	})
+}
+
+// WithProxy configures a file proxy using the given txtar-encoded string.
+func WithProxy(txt string) RunOption {
+	return optionSetter(func(opts *runConfig) {
+		opts.proxyTxt = txt
+	})
+}
+
+// WithModes configures the execution modes that the test should run in.
+func WithModes(modes EnvMode) RunOption {
+	return optionSetter(func(opts *runConfig) {
+		opts.modes = modes
+	})
+}
+
 // Run executes the test function in the default configured gopls execution
 // modes. For each a test run, a new workspace is created containing the
 // un-txtared files specified by filedata.
-func (r *Runner) Run(t *testing.T, filedata string, test func(t *testing.T, e *Env)) {
+func (r *Runner) Run(t *testing.T, filedata string, test func(t *testing.T, e *Env), opts ...RunOption) {
 	t.Helper()
-	r.RunInMode(r.defaultModes, t, filedata, test)
-}
+	config := r.defaultConfig()
+	for _, opt := range opts {
+		opt.set(config)
+	}
 
-// RunInMode runs the test in the execution modes specified by the modes bitmask.
-func (r *Runner) RunInMode(modes EnvMode, t *testing.T, filedata string, test func(t *testing.T, e *Env)) {
-	t.Helper()
 	tests := []struct {
 		name      string
 		mode      EnvMode
@@ -186,16 +230,16 @@ func (r *Runner) RunInMode(modes EnvMode, t *testing.T, filedata string, test fu
 
 	for _, tc := range tests {
 		tc := tc
-		if modes&tc.mode == 0 {
+		if config.modes&tc.mode == 0 {
 			continue
 		}
 		t.Run(tc.name, func(t *testing.T) {
 			t.Helper()
-			ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
+			ctx, cancel := context.WithTimeout(context.Background(), config.timeout)
 			defer cancel()
 			ctx = debug.WithInstance(ctx, "", "")
 
-			ws, err := fake.NewWorkspace("regtest", []byte(filedata))
+			ws, err := fake.NewWorkspace("regtest", filedata, config.proxyTxt)
 			if err != nil {
 				t.Fatal(err)
 			}
diff --git a/internal/lsp/regtest/shared_test.go b/internal/lsp/regtest/shared_test.go
index adc7283b1..0d53bbb74 100644
--- a/internal/lsp/regtest/shared_test.go
+++ b/internal/lsp/regtest/shared_test.go
@@ -25,12 +25,12 @@ func main() {
 func runShared(t *testing.T, program string, testFunc func(env1 *Env, env2 *Env)) {
 	// Only run these tests in forwarded modes.
 	modes := runner.Modes() & (Forwarded | SeparateProcess)
-	runner.RunInMode(modes, t, sharedProgram, func(t *testing.T, env1 *Env) {
+	runner.Run(t, sharedProgram, func(t *testing.T, env1 *Env) {
 		// Create a second test session connected to the same workspace and server
 		// as the first.
 		env2 := NewEnv(env1.Ctx, t, env1.W, env1.Server)
 		testFunc(env1, env2)
-	})
+	}, WithModes(modes))
 }
 
 func TestSimultaneousEdits(t *testing.T) {