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