diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go index baf5fbd..30ca84c 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 6b17127..d0617c2 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) + } + }) + } +}