1package main
2
3import (
4	"bytes"
5	"flag"
6	"go/build"
7	"go/parser"
8	"go/token"
9	"io/ioutil"
10	"os"
11	"os/exec"
12	"path/filepath"
13	"regexp"
14	"runtime"
15	"strings"
16	"testing"
17)
18
19// Set --regenerate to regenerate the golden files.
20var regenerate = flag.Bool("regenerate", false, "regenerate golden files")
21
22// When the environment variable RUN_AS_PROTOC_GEN_GO is set, we skip running
23// tests and instead act as protoc-gen-go. This allows the test binary to
24// pass itself to protoc.
25func init() {
26	if os.Getenv("RUN_AS_PROTOC_GEN_GO") != "" {
27		main()
28		os.Exit(0)
29	}
30}
31
32func TestGolden(t *testing.T) {
33	workdir, err := ioutil.TempDir("", "proto-test")
34	if err != nil {
35		t.Fatal(err)
36	}
37	defer os.RemoveAll(workdir)
38
39	// Find all the proto files we need to compile. We assume that each directory
40	// contains the files for a single package.
41	supportTypeAliases := hasReleaseTag("go1.9")
42	packages := map[string][]string{}
43	err = filepath.Walk("testdata", func(path string, info os.FileInfo, err error) error {
44		if filepath.Base(path) == "import_public" && !supportTypeAliases {
45			// Public imports require type alias support.
46			return filepath.SkipDir
47		}
48		if !strings.HasSuffix(path, ".proto") {
49			return nil
50		}
51		dir := filepath.Dir(path)
52		packages[dir] = append(packages[dir], path)
53		return nil
54	})
55	if err != nil {
56		t.Fatal(err)
57	}
58
59	// Compile each package, using this binary as protoc-gen-go.
60	for _, sources := range packages {
61		args := []string{"-Itestdata", "--go_out=plugins=grpc,paths=source_relative:" + workdir}
62		args = append(args, sources...)
63		protoc(t, args)
64	}
65
66	// Compare each generated file to the golden version.
67	filepath.Walk(workdir, func(genPath string, info os.FileInfo, _ error) error {
68		if info.IsDir() {
69			return nil
70		}
71
72		// For each generated file, figure out the path to the corresponding
73		// golden file in the testdata directory.
74		relPath, err := filepath.Rel(workdir, genPath)
75		if err != nil {
76			t.Errorf("filepath.Rel(%q, %q): %v", workdir, genPath, err)
77			return nil
78		}
79		if filepath.SplitList(relPath)[0] == ".." {
80			t.Errorf("generated file %q is not relative to %q", genPath, workdir)
81		}
82		goldenPath := filepath.Join("testdata", relPath)
83
84		got, err := ioutil.ReadFile(genPath)
85		if err != nil {
86			t.Error(err)
87			return nil
88		}
89		if *regenerate {
90			// If --regenerate set, just rewrite the golden files.
91			err := ioutil.WriteFile(goldenPath, got, 0666)
92			if err != nil {
93				t.Error(err)
94			}
95			return nil
96		}
97
98		want, err := ioutil.ReadFile(goldenPath)
99		if err != nil {
100			t.Error(err)
101			return nil
102		}
103
104		want = fdescRE.ReplaceAll(want, nil)
105		got = fdescRE.ReplaceAll(got, nil)
106		if bytes.Equal(got, want) {
107			return nil
108		}
109
110		cmd := exec.Command("diff", "-u", goldenPath, genPath)
111		out, _ := cmd.CombinedOutput()
112		t.Errorf("golden file differs: %v\n%v", relPath, string(out))
113		return nil
114	})
115}
116
117var fdescRE = regexp.MustCompile(`(?ms)^var fileDescriptor.*}`)
118
119// Source files used by TestParameters.
120const (
121	aProto = `
122syntax = "proto3";
123package test.alpha;
124option go_package = "package/alpha";
125import "beta/b.proto";
126message M { test.beta.M field = 1; }`
127
128	bProto = `
129syntax = "proto3";
130package test.beta;
131// no go_package option
132message M {}`
133)
134
135func TestParameters(t *testing.T) {
136	for _, test := range []struct {
137		parameters   string
138		wantFiles    map[string]bool
139		wantImportsA map[string]bool
140		wantPackageA string
141		wantPackageB string
142	}{{
143		parameters: "",
144		wantFiles: map[string]bool{
145			"package/alpha/a.pb.go": true,
146			"beta/b.pb.go":          true,
147		},
148		wantPackageA: "alpha",
149		wantPackageB: "test_beta",
150		wantImportsA: map[string]bool{
151			"github.com/golang/protobuf/proto": true,
152			"beta":                             true,
153		},
154	}, {
155		parameters: "import_prefix=prefix",
156		wantFiles: map[string]bool{
157			"package/alpha/a.pb.go": true,
158			"beta/b.pb.go":          true,
159		},
160		wantPackageA: "alpha",
161		wantPackageB: "test_beta",
162		wantImportsA: map[string]bool{
163			// This really doesn't seem like useful behavior.
164			"prefixgithub.com/golang/protobuf/proto": true,
165			"prefixbeta":                             true,
166		},
167	}, {
168		// import_path only affects the 'package' line.
169		parameters:   "import_path=import/path/of/pkg",
170		wantPackageA: "alpha",
171		wantPackageB: "pkg",
172		wantFiles: map[string]bool{
173			"package/alpha/a.pb.go": true,
174			"beta/b.pb.go":          true,
175		},
176	}, {
177		parameters: "Mbeta/b.proto=package/gamma",
178		wantFiles: map[string]bool{
179			"package/alpha/a.pb.go": true,
180			"beta/b.pb.go":          true,
181		},
182		wantPackageA: "alpha",
183		wantPackageB: "test_beta",
184		wantImportsA: map[string]bool{
185			"github.com/golang/protobuf/proto": true,
186			// Rewritten by the M parameter.
187			"package/gamma": true,
188		},
189	}, {
190		parameters: "import_prefix=prefix,Mbeta/b.proto=package/gamma",
191		wantFiles: map[string]bool{
192			"package/alpha/a.pb.go": true,
193			"beta/b.pb.go":          true,
194		},
195		wantPackageA: "alpha",
196		wantPackageB: "test_beta",
197		wantImportsA: map[string]bool{
198			// import_prefix applies after M.
199			"prefixpackage/gamma": true,
200		},
201	}, {
202		parameters: "paths=source_relative",
203		wantFiles: map[string]bool{
204			"alpha/a.pb.go": true,
205			"beta/b.pb.go":  true,
206		},
207		wantPackageA: "alpha",
208		wantPackageB: "test_beta",
209	}, {
210		parameters: "paths=source_relative,import_prefix=prefix",
211		wantFiles: map[string]bool{
212			// import_prefix doesn't affect filenames.
213			"alpha/a.pb.go": true,
214			"beta/b.pb.go":  true,
215		},
216		wantPackageA: "alpha",
217		wantPackageB: "test_beta",
218	}} {
219		name := test.parameters
220		if name == "" {
221			name = "defaults"
222		}
223		// TODO: Switch to t.Run when we no longer support Go 1.6.
224		t.Logf("TEST: %v", name)
225		workdir, err := ioutil.TempDir("", "proto-test")
226		if err != nil {
227			t.Fatal(err)
228		}
229		defer os.RemoveAll(workdir)
230
231		for _, dir := range []string{"alpha", "beta", "out"} {
232			if err := os.MkdirAll(filepath.Join(workdir, dir), 0777); err != nil {
233				t.Fatal(err)
234			}
235		}
236
237		if err := ioutil.WriteFile(filepath.Join(workdir, "alpha", "a.proto"), []byte(aProto), 0666); err != nil {
238			t.Fatal(err)
239		}
240
241		if err := ioutil.WriteFile(filepath.Join(workdir, "beta", "b.proto"), []byte(bProto), 0666); err != nil {
242			t.Fatal(err)
243		}
244
245		protoc(t, []string{
246			"-I" + workdir,
247			"--go_out=" + test.parameters + ":" + filepath.Join(workdir, "out"),
248			filepath.Join(workdir, "alpha", "a.proto"),
249		})
250		protoc(t, []string{
251			"-I" + workdir,
252			"--go_out=" + test.parameters + ":" + filepath.Join(workdir, "out"),
253			filepath.Join(workdir, "beta", "b.proto"),
254		})
255
256		contents := make(map[string]string)
257		gotFiles := make(map[string]bool)
258		outdir := filepath.Join(workdir, "out")
259		filepath.Walk(outdir, func(p string, info os.FileInfo, _ error) error {
260			if info.IsDir() {
261				return nil
262			}
263			base := filepath.Base(p)
264			if base == "a.pb.go" || base == "b.pb.go" {
265				b, err := ioutil.ReadFile(p)
266				if err != nil {
267					t.Fatal(err)
268				}
269				contents[base] = string(b)
270			}
271			relPath, _ := filepath.Rel(outdir, p)
272			gotFiles[relPath] = true
273			return nil
274		})
275		for got := range gotFiles {
276			if runtime.GOOS == "windows" {
277				got = filepath.ToSlash(got)
278			}
279			if !test.wantFiles[got] {
280				t.Errorf("unexpected output file: %v", got)
281			}
282		}
283		for want := range test.wantFiles {
284			if runtime.GOOS == "windows" {
285				want = filepath.FromSlash(want)
286			}
287			if !gotFiles[want] {
288				t.Errorf("missing output file:    %v", want)
289			}
290		}
291		gotPackageA, gotImports, err := parseFile(contents["a.pb.go"])
292		if err != nil {
293			t.Fatal(err)
294		}
295		gotPackageB, _, err := parseFile(contents["b.pb.go"])
296		if err != nil {
297			t.Fatal(err)
298		}
299		if got, want := gotPackageA, test.wantPackageA; want != got {
300			t.Errorf("output file a.pb.go is package %q, want %q", got, want)
301		}
302		if got, want := gotPackageB, test.wantPackageB; want != got {
303			t.Errorf("output file b.pb.go is package %q, want %q", got, want)
304		}
305		missingImport := false
306	WantImport:
307		for want := range test.wantImportsA {
308			for _, imp := range gotImports {
309				if `"`+want+`"` == imp {
310					continue WantImport
311				}
312			}
313			t.Errorf("output file a.pb.go does not contain expected import %q", want)
314			missingImport = true
315		}
316		if missingImport {
317			t.Error("got imports:")
318			for _, imp := range gotImports {
319				t.Errorf("  %v", imp)
320			}
321		}
322	}
323}
324
325// parseFile returns a file's package name and a list of all packages it imports.
326func parseFile(source string) (packageName string, imports []string, err error) {
327	fset := token.NewFileSet()
328	f, err := parser.ParseFile(fset, "<source>", source, parser.ImportsOnly)
329	if err != nil {
330		return "", nil, err
331	}
332	for _, imp := range f.Imports {
333		imports = append(imports, imp.Path.Value)
334	}
335	return f.Name.Name, imports, nil
336}
337
338func protoc(t *testing.T, args []string) {
339	cmd := exec.Command("protoc", "--plugin=protoc-gen-go="+os.Args[0])
340	cmd.Args = append(cmd.Args, args...)
341	// We set the RUN_AS_PROTOC_GEN_GO environment variable to indicate that
342	// the subprocess should act as a proto compiler rather than a test.
343	cmd.Env = append(os.Environ(), "RUN_AS_PROTOC_GEN_GO=1")
344	out, err := cmd.CombinedOutput()
345	if len(out) > 0 || err != nil {
346		t.Log("RUNNING: ", strings.Join(cmd.Args, " "))
347	}
348	if len(out) > 0 {
349		t.Log(string(out))
350	}
351	if err != nil {
352		t.Fatalf("protoc: %v", err)
353	}
354}
355
356func hasReleaseTag(want string) bool {
357	for _, tag := range build.Default.ReleaseTags {
358		if tag == want {
359			return true
360		}
361	}
362	return false
363}
364