Fix -write_generate_directive behavior for "go tool mockgen" invocations

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 <constraint>")
 	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 {

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)
+			}
+		})
+	}
+}