go/packages/packagestest: allow tests to also specify overlay contents

This is needed to write the more advanced test cases for editor integrated
tools.
The expectation system also uses the overlay rather than the file if it is
supplied.

Change-Id: I8fd21f4efe5ac5869fa6e25d3cd0d5096051e5e5
Reviewed-on: https://go-review.googlesource.com/c/153240
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/go/packages/packagestest/expect.go b/go/packages/packagestest/expect.go
index 14e1fc0..0902057 100644
--- a/go/packages/packagestest/expect.go
+++ b/go/packages/packagestest/expect.go
@@ -146,7 +146,11 @@
 	}
 	for _, pkg := range pkgs {
 		for _, filename := range pkg.GoFiles {
-			l, err := expect.Parse(e.fset, filename, nil)
+			content, err := e.FileContents(filename)
+			if err != nil {
+				return err
+			}
+			l, err := expect.Parse(e.fset, filename, content)
 			if err != nil {
 				return fmt.Errorf("Failed to extract expectations: %v", err)
 			}
@@ -337,7 +341,7 @@
 			return mark, args, nil
 		}
 	case string:
-		start, end, err := expect.MatchBefore(e.fset, e.fileContents, n.Pos, arg)
+		start, end, err := expect.MatchBefore(e.fset, e.FileContents, n.Pos, arg)
 		if err != nil {
 			return Range{}, nil, err
 		}
@@ -346,7 +350,7 @@
 		}
 		return Range{Start: start, End: end}, args, nil
 	case *regexp.Regexp:
-		start, end, err := expect.MatchBefore(e.fset, e.fileContents, n.Pos, arg)
+		start, end, err := expect.MatchBefore(e.fset, e.FileContents, n.Pos, arg)
 		if err != nil {
 			return Range{}, nil, err
 		}
diff --git a/go/packages/packagestest/export.go b/go/packages/packagestest/export.go
index 7214edd..14a0ac1 100644
--- a/go/packages/packagestest/export.go
+++ b/go/packages/packagestest/export.go
@@ -39,6 +39,11 @@
 	// be a string or byte slice, in which case it is the contents of the
 	// file, otherwise it must be a Writer function.
 	Files map[string]interface{}
+
+	// Overlay is the set of source file overlays for the module.
+	// The keys are the file fragment as in the Files configuration.
+	// The values are the in memory overlay content for the file.
+	Overlay map[string][]byte
 }
 
 // A Writer is a function that writes out a test file.
@@ -53,13 +58,15 @@
 	// Exactly what it will contain varies depending on the Exporter being used.
 	Config *packages.Config
 
-	temp     string                       // the temporary directory that was exported to
-	primary  string                       // the first non GOROOT module that was exported
-	written  map[string]map[string]string // the full set of exported files
-	fset     *token.FileSet               // The file set used when parsing expectations
-	notes    []*expect.Note               // The list of expectations extracted from go source files
-	markers  map[string]Range             // The set of markers extracted from go source files
-	contents map[string][]byte
+	// Modules is the module description that was used to produce this exported data set.
+	Modules []Module
+
+	temp    string                       // the temporary directory that was exported to
+	primary string                       // the first non GOROOT module that was exported
+	written map[string]map[string]string // the full set of exported files
+	fset    *token.FileSet               // The file set used when parsing expectations
+	notes   []*expect.Note               // The list of expectations extracted from go source files
+	markers map[string]Range             // The set of markers extracted from go source files
 }
 
 // Exporter implementations are responsible for converting from the generic description of some
@@ -126,14 +133,15 @@
 	}
 	exported := &Exported{
 		Config: &packages.Config{
-			Dir: temp,
-			Env: append(os.Environ(), "GOPACKAGESDRIVER=off"),
+			Dir:     temp,
+			Env:     append(os.Environ(), "GOPACKAGESDRIVER=off"),
+			Overlay: make(map[string][]byte),
 		},
-		temp:     temp,
-		primary:  modules[0].Name,
-		written:  map[string]map[string]string{},
-		fset:     token.NewFileSet(),
-		contents: map[string][]byte{},
+		Modules: modules,
+		temp:    temp,
+		primary: modules[0].Name,
+		written: map[string]map[string]string{},
+		fset:    token.NewFileSet(),
 	}
 	defer func() {
 		if t.Failed() || t.Skipped() {
@@ -165,6 +173,10 @@
 				t.Fatalf("Invalid type %T in files, must be string or Writer", value)
 			}
 		}
+		for fragment, value := range module.Overlay {
+			fullpath := exporter.Filename(exported, module.Name, filepath.FromSlash(fragment))
+			exported.Config.Overlay[fullpath] = value
+		}
 	}
 	if err := exporter.Finalize(exported); err != nil {
 		t.Fatal(err)
@@ -282,8 +294,11 @@
 	return ""
 }
 
-func (e *Exported) fileContents(filename string) ([]byte, error) {
-	if content, found := e.contents[filename]; found {
+// FileContents returns the contents of the specified file.
+// It will use the overlay if the file is present, otherwise it will read it
+// from disk.
+func (e *Exported) FileContents(filename string) ([]byte, error) {
+	if content, found := e.Config.Overlay[filename]; found {
 		return content, nil
 	}
 	content, err := ioutil.ReadFile(filename)
diff --git a/go/packages/packagestest/export_test.go b/go/packages/packagestest/export_test.go
index decfd7c..e2a2355 100644
--- a/go/packages/packagestest/export_test.go
+++ b/go/packages/packagestest/export_test.go
@@ -5,7 +5,6 @@
 package packagestest_test
 
 import (
-	"io/ioutil"
 	"os"
 	"path/filepath"
 	"testing"
@@ -17,7 +16,11 @@
 	Name: "golang.org/fake1",
 	Files: map[string]interface{}{
 		"a.go": packagestest.Symlink("testdata/a.go"),
-		"b.go": "package fake1",
+		"b.go": "invalid file contents",
+	},
+	Overlay: map[string][]byte{
+		"b.go": []byte("package fake1"),
+		"c.go": []byte("package fake1"),
 	},
 }, {
 	Name: "golang.org/fake2",
@@ -33,7 +36,7 @@
 
 type fileTest struct {
 	module, fragment, expect string
-	check                    func(t *testing.T, filename string)
+	check                    func(t *testing.T, exported *packagestest.Exported, filename string)
 }
 
 func checkFiles(t *testing.T, exported *packagestest.Exported, tests []fileTest) {
@@ -46,14 +49,14 @@
 			t.Errorf("Got file %v, expected %v", got, expect)
 		}
 		if test.check != nil {
-			test.check(t, got)
+			test.check(t, exported, got)
 		}
 	}
 }
 
-func checkLink(expect string) func(t *testing.T, filename string) {
+func checkLink(expect string) func(t *testing.T, exported *packagestest.Exported, filename string) {
 	expect = filepath.FromSlash(expect)
-	return func(t *testing.T, filename string) {
+	return func(t *testing.T, exported *packagestest.Exported, filename string) {
 		if target, err := os.Readlink(filename); err != nil {
 			t.Errorf("Error checking link %v: %v", filename, err)
 		} else if target != expect {
@@ -62,9 +65,9 @@
 	}
 }
 
-func checkContent(expect string) func(t *testing.T, filename string) {
-	return func(t *testing.T, filename string) {
-		if content, err := ioutil.ReadFile(filename); err != nil {
+func checkContent(expect string) func(t *testing.T, exported *packagestest.Exported, filename string) {
+	return func(t *testing.T, exported *packagestest.Exported, filename string) {
+		if content, err := exported.FileContents(filename); err != nil {
 			t.Errorf("Error reading %v: %v", filename, err)
 		} else if string(content) != expect {
 			t.Errorf("Content of %v does not match, got %v expected %v", filename, string(content), expect)