go/packages/packagestest: fix grouping to account for module versions

This change adds support within packagestest for multiple versions of a supporting module. An example folder structure could now be:

- primarymod
- modules
  - repo1
    - mod1@v1.0.0
      - files
    - mod1@v1.1.0
    - mod2

Since packagestest creates a fake module proxy, we can now supply that proxy with multiple versions of a single module. This will allow us to add more comprehension to tests and be able to test new features (i.e: CodeLens for upgrading dependencies).

Fixes golang/go#36091

Change-Id: I7516e2525228e78b3f5bf1f7d1da1e42038dee37
Reviewed-on: https://go-review.googlesource.com/c/tools/+/218502
Run-TryBot: Rohan Challa <rohan@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

10 files changed, 192 insertions, 88 deletions

diff --git a/go/packages/packagestest/export.go b/go/packages/packagestest/export.go
index d6fd0ab19..d702cb4b1 100644
--- a/go/packages/packagestest/export.go
+++ b/go/packages/packagestest/export.go
@@ -364,13 +364,38 @@ func GroupFilesByModules(root string) ([]Module, error) {
 	}
 
 	var currentRepo, currentModule string
+	updateCurrentModule := func(dir string) {
+		if dir == currentModule {
+			return
+		}
+		// Handle the case where we step into a nested directory that is a module
+		// and then step out into the parent which is also a module.
+		// Example:
+		// - repoa
+		//   - moda
+		//     - go.mod
+		//     - v2
+		//       - go.mod
+		//     - what.go
+		//   - modb
+		for dir != root {
+			if mods[dir] != nil {
+				currentModule = dir
+				return
+			}
+			dir = filepath.Dir(dir)
+		}
+	}
+
 	if err := filepath.Walk(modulesPath, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
 			return err
 		}
+		enclosingDir := filepath.Dir(path)
 		// If the path is not a directory, then we want to add the path to
 		// the files map of the currentModule.
 		if !info.IsDir() {
+			updateCurrentModule(enclosingDir)
 			fragment, err := filepath.Rel(currentModule, path)
 			if err != nil {
 				return err
@@ -380,13 +405,14 @@ func GroupFilesByModules(root string) ([]Module, error) {
 		}
 		// If the path is a directory and it's enclosing folder is equal to
 		// the modules folder, then the path is a new repo.
-		if filepath.Dir(path) == modulesPath {
+		if enclosingDir == modulesPath {
 			currentRepo = path
 			return nil
 		}
 		// If the path is a directory and it's enclosing folder is not the same
-		// as the current repo, then the path is a folder/package of the current module.
-		if filepath.Dir(path) != currentRepo {
+		// as the current repo and it is not of the form `v1`,`v2`,...
+		// then the path is a folder/package of the current module.
+		if enclosingDir != currentRepo && !versionSuffixRE.MatchString(filepath.Base(path)) {
 			return nil
 		}
 		// If the path is a directory and it's enclosing folder is the current repo
diff --git a/go/packages/packagestest/export_test.go b/go/packages/packagestest/export_test.go
index e2a2355fa..e7bf124c4 100644
--- a/go/packages/packagestest/export_test.go
+++ b/go/packages/packagestest/export_test.go
@@ -32,6 +32,16 @@ var testdata = []packagestest.Module{{
 	Files: map[string]interface{}{
 		"other/a.go": "package fake2",
 	},
+}, {
+	Name: "golang.org/fake3@v1.0.0",
+	Files: map[string]interface{}{
+		"other/a.go": "package fake3",
+	},
+}, {
+	Name: "golang.org/fake3@v1.1.0",
+	Files: map[string]interface{}{
+		"other/a.go": "package fake3",
+	},
 }}
 
 type fileTest struct {
@@ -74,3 +84,98 @@ func checkContent(expect string) func(t *testing.T, exported *packagestest.Expor
 		}
 	}
 }
+
+func TestGroupFilesByModules(t *testing.T) {
+	for _, tt := range []struct {
+		testdir string
+		want    []packagestest.Module
+	}{
+		{
+			testdir: "testdata/groups/one",
+			want: []packagestest.Module{
+				{
+					Name: "testdata/groups/one",
+					Files: map[string]interface{}{
+						"main.go": true,
+					},
+				},
+				{
+					Name: "example.com/extra",
+					Files: map[string]interface{}{
+						"help.go": true,
+					},
+				},
+			},
+		},
+		{
+			testdir: "testdata/groups/two",
+			want: []packagestest.Module{
+				{
+					Name: "testdata/groups/two",
+					Files: map[string]interface{}{
+						"main.go":      true,
+						"expect/yo.go": true,
+					},
+				},
+				{
+					Name: "example.com/extra",
+					Files: map[string]interface{}{
+						"yo.go":        true,
+						"geez/help.go": true,
+					},
+				},
+				{
+					Name: "example.com/extra/v2",
+					Files: map[string]interface{}{
+						"me.go":        true,
+						"geez/help.go": true,
+					},
+				},
+				{
+					Name: "example.com/tempmod",
+					Files: map[string]interface{}{
+						"main.go": true,
+					},
+				},
+				{
+					Name: "example.com/what@v1.0.0",
+					Files: map[string]interface{}{
+						"main.go": true,
+					},
+				},
+				{
+					Name: "example.com/what@v1.1.0",
+					Files: map[string]interface{}{
+						"main.go": true,
+					},
+				},
+			},
+		},
+	} {
+		t.Run(tt.testdir, func(t *testing.T) {
+			got, err := packagestest.GroupFilesByModules(tt.testdir)
+			if err != nil {
+				t.Fatalf("could not group files %v", err)
+			}
+			if len(got) != len(tt.want) {
+				t.Fatalf("%s: wanted %d modules but got %d", tt.testdir, len(tt.want), len(got))
+			}
+			for i, w := range tt.want {
+				g := got[i]
+				if filepath.FromSlash(g.Name) != filepath.FromSlash(w.Name) {
+					t.Fatalf("%s: wanted module[%d].Name to be %s but got %s", tt.testdir, i, filepath.FromSlash(w.Name), filepath.FromSlash(g.Name))
+				}
+				for fh := range w.Files {
+					if _, ok := g.Files[fh]; !ok {
+						t.Fatalf("%s, module[%d]: wanted %s but could not find", tt.testdir, i, fh)
+					}
+				}
+				for fh := range g.Files {
+					if _, ok := w.Files[fh]; !ok {
+						t.Fatalf("%s, module[%d]: found unexpected file %s", tt.testdir, i, fh)
+					}
+				}
+			}
+		})
+	}
+}
diff --git a/go/packages/packagestest/gopath_test.go b/go/packages/packagestest/gopath_test.go
index d75d4ce04..a251dbce2 100644
--- a/go/packages/packagestest/gopath_test.go
+++ b/go/packages/packagestest/gopath_test.go
@@ -26,79 +26,3 @@ func TestGOPATHExport(t *testing.T) {
 		{"golang.org/fake2/v2", "other/a.go", "fake2_v2/src/golang.org/fake2/v2/other/a.go", checkContent("package fake2")},
 	})
 }
-
-func TestGroupFilesByModules(t *testing.T) {
-	for _, tt := range []struct {
-		testdir string
-		want    []packagestest.Module
-	}{
-		{
-			testdir: "testdata/groups/one",
-			want: []packagestest.Module{
-				{
-					Name: "testdata/groups/one",
-					Files: map[string]interface{}{
-						"main.go": true,
-					},
-				},
-				{
-					Name: "example.com/extra",
-					Files: map[string]interface{}{
-						"help.go": true,
-					},
-				},
-			},
-		},
-		{
-			testdir: "testdata/groups/two",
-			want: []packagestest.Module{
-				{
-					Name: "testdata/groups/two",
-					Files: map[string]interface{}{
-						"main.go":      true,
-						"expect/yo.go": true,
-					},
-				},
-				{
-					Name: "example.com/extra",
-					Files: map[string]interface{}{
-						"me.go":        true,
-						"geez/help.go": true,
-					},
-				},
-				{
-					Name: "example.com/tempmod",
-					Files: map[string]interface{}{
-						"main.go": true,
-					},
-				},
-			},
-		},
-	} {
-		t.Run(tt.testdir, func(t *testing.T) {
-			got, err := packagestest.GroupFilesByModules(tt.testdir)
-			if err != nil {
-				t.Fatalf("could not group files %v", err)
-			}
-			if len(got) != len(tt.want) {
-				t.Fatalf("%s: wanted %d modules but got %d", tt.testdir, len(tt.want), len(got))
-			}
-			for i, w := range tt.want {
-				g := got[i]
-				if filepath.FromSlash(g.Name) != filepath.FromSlash(w.Name) {
-					t.Fatalf("%s: wanted module[%d].Name to be %s but got %s", tt.testdir, i, filepath.FromSlash(w.Name), filepath.FromSlash(g.Name))
-				}
-				for fh := range w.Files {
-					if _, ok := g.Files[fh]; !ok {
-						t.Fatalf("%s, module[%d]: wanted %s but could not find", tt.testdir, i, fh)
-					}
-				}
-				for fh := range g.Files {
-					if _, ok := w.Files[fh]; !ok {
-						t.Fatalf("%s, module[%d]: found unexpected file %s", tt.testdir, i, fh)
-					}
-				}
-			}
-		})
-	}
-}
diff --git a/go/packages/packagestest/modules.go b/go/packages/packagestest/modules.go
index 44712043b..e393cdca0 100644
--- a/go/packages/packagestest/modules.go
+++ b/go/packages/packagestest/modules.go
@@ -14,6 +14,7 @@ import (
 	"path"
 	"path/filepath"
 	"regexp"
+	"strings"
 
 	"golang.org/x/tools/go/packages"
 )
@@ -40,6 +41,11 @@ var Modules = modules{}
 
 type modules struct{}
 
+type moduleAtVersion struct {
+	module  string
+	version string
+}
+
 func (modules) Name() string {
 	return "Modules"
 }
@@ -64,6 +70,18 @@ func (modules) Finalize(exported *Exported) error {
 		exported.written[exported.primary] = make(map[string]string)
 	}
 
+	// Create a map of modulepath -> {module, version} for modulepaths
+	// that are of the form `repoa/mod1@v1.1.0`.
+	versions := make(map[string]moduleAtVersion)
+	for module := range exported.written {
+		if splt := strings.Split(module, "@"); len(splt) > 1 {
+			versions[module] = moduleAtVersion{
+				module:  splt[0],
+				version: splt[1],
+			}
+		}
+	}
+
 	// If the primary module already has a go.mod, write the contents to a temp
 	// go.mod for now and then we will reset it when we are getting all the markers.
 	if gomod := exported.written[exported.primary]["go.mod"]; gomod != "" {
@@ -82,7 +100,14 @@ func (modules) Finalize(exported *Exported) error {
 		if other == exported.primary {
 			continue
 		}
-		primaryGomod += fmt.Sprintf("\t%v %v\n", other, moduleVersion(other))
+		version := moduleVersion(other)
+		// If other is of the form `repo1/mod1@v1.1.0`,
+		// then we need to extract the module and the version.
+		if v, ok := versions[other]; ok {
+			other = v.module
+			version = v.version
+		}
+		primaryGomod += fmt.Sprintf("\t%v %v\n", other, version)
 	}
 	primaryGomod += ")\n"
 	if err := ioutil.WriteFile(filepath.Join(primaryDir, "go.mod"), []byte(primaryGomod), 0644); err != nil {
@@ -100,8 +125,12 @@ func (modules) Finalize(exported *Exported) error {
 			continue
 		}
 		dir := moduleDir(exported, module)
-
 		modfile := filepath.Join(dir, "go.mod")
+		// If other is of the form `repo1/mod1@v1.1.0`,
+		// then we need to extract the module name without the version.
+		if v, ok := versions[module]; ok {
+			module = v.module
+		}
 		if err := ioutil.WriteFile(modfile, []byte("module "+module+"\n"), 0644); err != nil {
 			return err
 		}
@@ -114,9 +143,15 @@ func (modules) Finalize(exported *Exported) error {
 		if module == exported.primary {
 			continue
 		}
+		version := moduleVersion(module)
+		// If other is of the form `repo1/mod1@v1.1.0`,
+		// then we need to extract the module and the version.
+		if v, ok := versions[module]; ok {
+			module = v.module
+			version = v.version
+		}
 		dir := filepath.Join(proxyDir, module, "@v")
-
-		if err := writeModuleProxy(dir, module, files); err != nil {
+		if err := writeModuleProxy(dir, module, version, files); err != nil {
 			return fmt.Errorf("creating module proxy dir for %v: %v", module, err)
 		}
 	}
@@ -142,14 +177,19 @@ func (modules) Finalize(exported *Exported) error {
 }
 
 // writeModuleProxy creates a directory in the proxy dir for a module.
-func writeModuleProxy(dir, module string, files map[string]string) error {
-	ver := moduleVersion(module)
+func writeModuleProxy(dir, module, ver string, files map[string]string) error {
 	if err := os.MkdirAll(dir, 0755); err != nil {
 		return err
 	}
 
-	// list file. Just the single version.
-	if err := ioutil.WriteFile(filepath.Join(dir, "list"), []byte(ver+"\n"), 0644); err != nil {
+	// the modproxy checks for versions by looking at the "list" file,
+	// since we are supporting multiple versions, create the file if it does not exist or
+	// append the version number to the preexisting file.
+	f, err := os.OpenFile(filepath.Join(dir, "list"), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		return err
+	}
+	if _, err := f.WriteString(ver + "\n"); err != nil {
 		return err
 	}
 
@@ -169,7 +209,7 @@ func writeModuleProxy(dir, module string, files map[string]string) error {
 	}
 
 	// zip of all the source files.
-	f, err := os.OpenFile(filepath.Join(dir, ver+".zip"), os.O_CREATE|os.O_WRONLY, 0644)
+	f, err = os.OpenFile(filepath.Join(dir, ver+".zip"), os.O_CREATE|os.O_WRONLY, 0644)
 	if err != nil {
 		return err
 	}
@@ -220,6 +260,9 @@ func primaryDir(exported *Exported) string {
 }
 
 func moduleDir(exported *Exported, module string) string {
+	if strings.Contains(module, "@") {
+		return filepath.Join(modCache(exported), module)
+	}
 	return filepath.Join(modCache(exported), path.Dir(module), path.Base(module)+"@"+moduleVersion(module))
 }
 
diff --git a/go/packages/packagestest/modules_test.go b/go/packages/packagestest/modules_test.go
index 0e88c4539..5d13176da 100644
--- a/go/packages/packagestest/modules_test.go
+++ b/go/packages/packagestest/modules_test.go
@@ -28,5 +28,7 @@ func TestModulesExport(t *testing.T) {
 		{"golang.org/fake2", "go.mod", "modcache/pkg/mod/golang.org/fake2@v1.0.0/go.mod", nil},
 		{"golang.org/fake2", "other/a.go", "modcache/pkg/mod/golang.org/fake2@v1.0.0/other/a.go", checkContent("package fake2")},
 		{"golang.org/fake2/v2", "other/a.go", "modcache/pkg/mod/golang.org/fake2/v2@v2.0.0/other/a.go", checkContent("package fake2")},
+		{"golang.org/fake3@v1.1.0", "other/a.go", "modcache/pkg/mod/golang.org/fake3@v1.1.0/other/a.go", checkContent("package fake3")},
+		{"golang.org/fake3@v1.0.0", "other/a.go", "modcache/pkg/mod/golang.org/fake3@v1.0.0/other/a.go", nil},
 	})
 }
diff --git a/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/v2/geez/help.go b/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/v2/geez/help.go
new file mode 100644
index 000000000..930ffdc81
--- /dev/null
+++ b/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/v2/geez/help.go
@@ -0,0 +1 @@
+package example.com/extra/geez
\ No newline at end of file
diff --git a/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/me.go b/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/v2/me.go
index 6a8c7d31f..6a8c7d31f 100644
--- a/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/me.go
+++ b/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/v2/me.go
diff --git a/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/yo.go b/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/yo.go
new file mode 100644
index 000000000..6a8c7d31f
--- /dev/null
+++ b/go/packages/packagestest/testdata/groups/two/modules/example.com/extra/yo.go
@@ -0,0 +1 @@
+package example.com/extra
\ No newline at end of file
diff --git a/go/packages/packagestest/testdata/groups/two/modules/example.com/what@v1.0.0/main.go b/go/packages/packagestest/testdata/groups/two/modules/example.com/what@v1.0.0/main.go
new file mode 100644
index 000000000..4723ee64b
--- /dev/null
+++ b/go/packages/packagestest/testdata/groups/two/modules/example.com/what@v1.0.0/main.go
@@ -0,0 +1 @@
+package example.com/what
\ No newline at end of file
diff --git a/go/packages/packagestest/testdata/groups/two/modules/example.com/what@v1.1.0/main.go b/go/packages/packagestest/testdata/groups/two/modules/example.com/what@v1.1.0/main.go
new file mode 100644
index 000000000..4723ee64b
--- /dev/null
+++ b/go/packages/packagestest/testdata/groups/two/modules/example.com/what@v1.1.0/main.go
@@ -0,0 +1 @@
+package example.com/what
\ No newline at end of file