From e5bd0603266966bb1d714f2e833da8d6dc62890c Mon Sep 17 00:00:00 2001 From: Scott Leggett Date: Thu, 10 Apr 2025 21:08:43 +0800 Subject: [PATCH] fix: support go tool invocation in write directive "go tool" invocations use dynamically built binaries to execute commands. This was not compatible with the way mockgen dynamically built //go:generate lines for -write_generate_directive, because os.Args[0] is a different temporary, dynamic value for each invocation of "go tool". The outcome was that this style of //go:generate command didn't work: //go:generate go tool mockgen -destination ./mock_handler_test.go -write_generate_directive ... Fix this by checking if mockgen is invoked via //go:generate. In that case, copy the //go:generate line from the source file instead of dynamically building it. Also add a -prog_name flag to override the detection of program name detection via arg[0]. This facilitates manual invocations of mockgen like: go tool mockgen -prog_name="go tool mockgen" -write_generate_directive .... --- mockgen/mockgen.go | 65 +++++++++++++++++++++++++---- mockgen/mockgen_test.go | 92 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+), 7 deletions(-) diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go index baf5fbdc..30ca84cd 100644 --- a/mockgen/mockgen.go +++ b/mockgen/mockgen.go @@ -19,6 +19,7 @@ package main // TODO: This does not support embedding package-local interfaces in a separate file. import ( + "bufio" "bytes" "encoding/json" "errors" @@ -64,6 +65,7 @@ var ( writePkgComment = flag.Bool("write_package_comment", true, "Writes package documentation comment (godoc) if true.") writeSourceComment = flag.Bool("write_source_comment", true, "Writes original file (source mode) or interface names (package mode) comment if true.") writeGenerateDirective = flag.Bool("write_generate_directive", false, "Add //go:generate directive to regenerate the mock") + progName = flag.String("prog_name", "", "The program name to use in the generate directive and command comment") copyrightFile = flag.String("copyright_file", "", "Copyright file used to add copyright header") buildConstraint = flag.String("build_constraint", "", "If non-empty, added as //go:build ") typed = flag.Bool("typed", false, "Generate Type-safe 'Return', 'Do', 'DoAndReturn' function") @@ -340,6 +342,60 @@ func sanitize(s string) string { return t } +// getGenerateDirective reconstructs the //go:generate directive used to invoke +// mockgen. If invoked via "go generate", it reads the exact line from the +// source file identified by gofile and goline (the GOFILE and GOLINE values +// set by "go generate"). This preserves the exact command prefix (e.g., +// "go tool mockgen"). If not invoked via "go generate" or the source file +// cannot be read, it falls back to reconstructing the command from args. If +// progName is set, it is used to override the program name (args[0]) in the +// fallback case. +func getGenerateDirective(progName, gofile, goline string, args []string) string { + cmdArgs := make([]string, len(args)) + copy(cmdArgs, args) + if progName != "" { + cmdArgs[0] = progName + } else { + name := filepath.Base(cmdArgs[0]) + if runtime.GOOS == "windows" { + name = strings.TrimSuffix(name, ".exe") + } + cmdArgs[0] = name + } + for i, arg := range cmdArgs { + if i > 0 && strings.ContainsAny(arg, " \t") { + cmdArgs[i] = strconv.Quote(arg) + } + } + fallback := "//go:generate " + strings.Join(cmdArgs, " ") + if gofile == "" || goline == "" { + return fallback + } + line, err := strconv.Atoi(goline) + if err != nil || line <= 0 { + return fallback + } + f, err := os.Open(gofile) + if err != nil { + return fallback + } + defer f.Close() + + scanner := bufio.NewScanner(f) + currentLine := 1 + for scanner.Scan() { + if currentLine == line { + genLine := strings.TrimSpace(scanner.Text()) + if strings.HasPrefix(genLine, "//go:generate ") { + return genLine + } + return fallback + } + currentLine++ + } + return fallback +} + func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPackagePath string) error { if outputPkgName != pkg.Name && *selfPackage == "" { // reset outputPackagePath if it's not passed in through -self_package @@ -372,12 +428,7 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac g.p("//") g.p("// Generated by this command:") g.p("//") - // only log the name of the executable, not the full path - name := filepath.Base(os.Args[0]) - if runtime.GOOS == "windows" { - name = strings.TrimSuffix(name, ".exe") - } - g.p("//\t%v", strings.Join(append([]string{name}, os.Args[1:]...), " ")) + g.p("//\t%v", strings.TrimPrefix(getGenerateDirective(*progName, os.Getenv("GOFILE"), os.Getenv("GOLINE"), os.Args), "//go:generate ")) g.p("//") } @@ -475,7 +526,7 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac g.p(")") if *writeGenerateDirective { - g.p("//go:generate %v", strings.Join(os.Args, " ")) + g.p("%v", getGenerateDirective(*progName, os.Getenv("GOFILE"), os.Getenv("GOLINE"), os.Args)) } for _, intf := range pkg.Interfaces { diff --git a/mockgen/mockgen_test.go b/mockgen/mockgen_test.go index 6b171272..d0617c2d 100644 --- a/mockgen/mockgen_test.go +++ b/mockgen/mockgen_test.go @@ -467,3 +467,95 @@ func TestParseExcludeInterfaces(t *testing.T) { }) } } + +func TestGetGenerateDirective(t *testing.T) { + // Setup a temporary file that acts as our GOFILE + tmpDir := t.TempDir() + tmpFile, err := os.Create(filepath.Join(tmpDir, "gofile.go")) + if err != nil { + t.Fatalf("failed to create temp file: %v", err) + } + + content := `package main + +//go:generate mockgen -destination=mock.go . MyInterface +//go:generate go tool mockgen -destination=mock2.go . MyOtherInterface +//go:generate mockgen -destination=mock3.go . SpacedInterface + //go:generate mockgen indented +func main() {} +` + if _, err := tmpFile.WriteString(content); err != nil { + t.Fatalf("failed to write to temp file: %v", err) + } + tmpFile.Close() + + osArgs := []string{"mockgen", "-destination=mock.go", ".", "MyInterface"} + + testCases := []struct { + name string + gofile string + goline string + progName string + expected string + }{ + { + name: "no env vars", + gofile: "", + goline: "", + expected: "//go:generate mockgen -destination=mock.go . MyInterface", + }, + { + name: "with progName", + gofile: "", + goline: "", + progName: "go tool mockgen", + expected: "//go:generate go tool mockgen -destination=mock.go . MyInterface", + }, + { + name: "line 1 (not a generate directive)", + gofile: tmpFile.Name(), + goline: "1", + expected: "//go:generate mockgen -destination=mock.go . MyInterface", + }, + { + name: "line 3 (standard mockgen)", + gofile: tmpFile.Name(), + goline: "3", + expected: "//go:generate mockgen -destination=mock.go . MyInterface", + }, + { + name: "line 4 (go tool mockgen)", + gofile: tmpFile.Name(), + goline: "4", + expected: "//go:generate go tool mockgen -destination=mock2.go . MyOtherInterface", + }, + { + name: "line 5 (spaced)", + gofile: tmpFile.Name(), + goline: "5", + expected: "//go:generate mockgen -destination=mock3.go . SpacedInterface", + }, + { + name: "line 6 (indented)", + gofile: tmpFile.Name(), + goline: "6", + expected: "//go:generate mockgen indented", + }, + { + name: "out of bounds line", + gofile: tmpFile.Name(), + goline: "100", + expected: "//go:generate mockgen -destination=mock.go . MyInterface", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + actual := getGenerateDirective(tc.progName, tc.gofile, tc.goline, osArgs) + if actual != tc.expected { + t.Errorf("expected %q, got %q", tc.expected, actual) + } + }) + } +}