add replace-type parameter

cmd/mockery.go

@@ -76,6 +76,7 @@ func NewRootCmd() *cobra.Command {
 	pFlags.Bool("unroll-variadic", true, "For functions with variadic arguments, do not unroll the arguments into the underlying testify call. Instead, pass variadic slice as-is.")
 	pFlags.Bool("exported", false, "Generates public mocks for private interfaces.")
 	pFlags.Bool("with-expecter", false, "Generate expecter utility around mock's On, Run and Return methods with explicit types. This option is NOT compatible with -unroll-variadic=false")
+	pFlags.StringArray("replace-type", nil, "Replace types")
 
 	viper.BindPFlags(pFlags)
 

cmd/mockery_test.go

@@ -47,6 +47,7 @@ func TestConfigEnvFlags(t *testing.T) {
 		UnrollVariadic:       false,
 		Exported:             true,
 		WithExpecter:         true,
+		ReplaceType:          []string{},
 	}
 
 	env(t, "CONFIG", expected.Config)
@@ -77,6 +78,7 @@ func TestConfigEnvFlags(t *testing.T) {
 	env(t, "UNROLL_VARIADIC", fmt.Sprint(expected.UnrollVariadic))
 	env(t, "EXPORTED", fmt.Sprint(expected.Exported))
 	env(t, "WITH_EXPECTER", fmt.Sprint(expected.WithExpecter))
+	env(t, "REPLACE_TYPE", strings.Join(expected.ReplaceType, ","))
 
 	initConfig(nil, nil)
 

docs/configuration.md

@@ -46,4 +46,5 @@ Parameter Descriptions
 | `exported` | Use `exported: True` to generate public mocks for private interfaces. |
 | `with-expecter` | Use `with-expecter: True` to generate `EXPECT()` methods for your mocks. This is the preferred way to setup your mocks. |
 | `testonly` | Prepend every mock file with `_test.go`. This is useful in cases where you are generating mocks `inpackage` but don't want the mocks to be visible to code outside of tests. |
-| `inpackage-suffix` | When `inpackage-suffix` is set to `True`, mock files are suffixed with `_mock` instead of being prefixed with `mock_` for InPackage mocks |
+| `inpackage-suffix` | When `inpackage-suffix` is set to `True`, mock files are suffixed with `_mock` instead of being prefixed with `mock_` for InPackage mocks |
+| `replace-type source=destination` | Replaces aliases, packages and/or types during generation.|

docs/features.md

@@ -68,3 +68,84 @@ Return(
     },
 )
 ```
+
+Replace Types
+-------------
+
+The `replace-type` parameter allows adding a list of type replacements to be made in package and/or type names.
+This can help overcome some parsing problems like type aliases that the Go parser doesn't provide enough information.
+
+```shell
+mockery --replace-type github.com/vektra/mockery/v2/baz/internal/foo.InternalBaz=baz:github.com/vektra/mockery/v2/baz.Baz
+```
+
+This parameter can be specified multiple times.
+
+This will replace any imported named `"github.com/vektra/mockery/v2/baz/internal/foo"`
+with `baz "github.com/vektra/mockery/v2/baz"`. The alias is defined with `:` before
+the package name. Also, the `InternalBaz` type that comes from this package will be renamed to `baz.Baz`.
+
+This next example fixes a common problem of type aliases that point to an internal package.
+
+`cloud.google.com/go/pubsub.Message` is a type alias defined like this:
+
+```go
+import (
+    ipubsub "cloud.google.com/go/internal/pubsub"
+)
+
+type Message = ipubsub.Message
+```
+
+The Go parser that mockery uses doesn't provide a way to detect this alias and sends the application the package and
+type name of the type in the internal package, which will not work.
+
+We can use "replace-type" with only the package part to replace any import of `cloud.google.com/go/internal/pubsub` to
+`cloud.google.com/go/pubsub`. We don't need to change the alias or type name in this case, because they are `pubsub`
+and `Message` in both cases.
+
+```shell
+mockery --replace-type cloud.google.com/go/internal/pubsub=cloud.google.com/go/pubsub
+```
+
+Original source:
+
+```go
+import (
+    "cloud.google.com/go/pubsub"
+)
+
+type Handler struct {
+    HandleMessage(m pubsub.Message) error
+}
+```
+
+Mock generated without this parameter:
+
+```go
+import (
+    mock "github.com/stretchr/testify/mock"
+
+    pubsub "cloud.google.com/go/internal/pubsub"
+)
+
+func (_m *Handler) HandleMessage(m pubsub.Message) error {
+    // ...
+    return nil
+}
+```
+
+Mock generated with this parameter.
+
+```go
+import (
+    mock "github.com/stretchr/testify/mock"
+
+    pubsub "cloud.google.com/go/pubsub"
+)
+
+func (_m *Handler) HandleMessage(m pubsub.Message) error {
+    // ...
+    return nil
+}
+```

pkg/config/config.go

@@ -52,5 +52,6 @@ type Config struct {
 	TestOnly       bool
 	UnrollVariadic bool `mapstructure:"unroll-variadic"`
 	Version        bool
-	WithExpecter   bool `mapstructure:"with-expecter"`
+	WithExpecter   bool     `mapstructure:"with-expecter"`
+	ReplaceType    []string `mapstructure:"replace-type"`
 }

pkg/fixtures/example_project/baz/foo.go

@@ -0,0 +1,12 @@
+package baz
+
+import (
+	ifoo "github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz/internal/foo"
+)
+
+type Baz = ifoo.InternalBaz
+
+type Foo interface {
+	DoFoo() string
+	GetBaz() (*Baz, error)
+}

pkg/fixtures/example_project/baz/internal/foo/foo.go

@@ -0,0 +1,6 @@
+package foo
+
+type InternalBaz struct {
+	One string
+	Two int
+}

pkg/generator.go

@@ -101,14 +101,48 @@ func (g *Generator) getPackageScopedType(ctx context.Context, o *types.TypeName)
 	if o.Pkg() == nil || o.Pkg().Name() == "main" || (!g.KeepTree && g.InPackage && o.Pkg() == g.iface.Pkg) {
 		return o.Name()
 	}
-	return g.addPackageImport(ctx, o.Pkg()) + "." + o.Name()
+	pkg := g.addPackageImport(ctx, o.Pkg())
+	name := o.Name()
+	g.checkReplaceType(ctx, func(from replaceType, to replaceType) bool {
+		if o.Pkg().Path() == from.pkg && name == from.typ {
+			name = to.typ
+			return false
+		}
+		return true
+	})
+	return pkg + "." + name
 }
 
 func (g *Generator) addPackageImport(ctx context.Context, pkg *types.Package) string {
 	return g.addPackageImportWithName(ctx, pkg.Path(), pkg.Name())
 }
 
+func (g *Generator) checkReplaceType(ctx context.Context, f func(from replaceType, to replaceType) bool) {
+	for _, replace := range g.ReplaceType {
+		r := strings.SplitN(replace, "=", 2)
+		if len(r) == 2 {
+			if !f(parseReplaceType(r[0]), parseReplaceType(r[1])) {
+				break
+			}
+		} else {
+			log := zerolog.Ctx(ctx)
+			log.Error().Msgf("invalid replace type value: %s", replace)
+		}
+	}
+}
+
 func (g *Generator) addPackageImportWithName(ctx context.Context, path, name string) string {
+	g.checkReplaceType(ctx, func(from replaceType, to replaceType) bool {
+		if path == from.pkg {
+			path = to.pkg
+			if to.alias != "" {
+				name = to.alias
+			}
+			return false
+		}
+		return true
+	})
+
 	if existingName, pathExists := g.packagePathToName[path]; pathExists {
 		return existingName
 	}
@@ -908,3 +942,27 @@ func resolveCollision(names []string, variable string) string {
 
 	return ret
 }
+
+type replaceType struct {
+	alias string
+	pkg   string
+	typ   string
+}
+
+func parseReplaceType(t string) replaceType {
+	ret := replaceType{}
+	r := strings.SplitN(t, ":", 2)
+	if len(r) > 1 {
+		ret.alias = r[0]
+		t = r[1]
+	}
+	lastDot := strings.LastIndex(t, ".")
+	lastSlash := strings.LastIndex(t, "/")
+	if lastDot == -1 || (lastSlash > -1 && lastDot < lastSlash) {
+		ret.pkg = t
+	} else {
+		ret.pkg = t[:lastDot]
+		ret.typ = t[lastDot+1:]
+	}
+	return ret
+}

pkg/generator_test.go

@@ -2411,6 +2411,90 @@ import mock "github.com/stretchr/testify/mock"
 	s.checkPrologueGeneration(generator, expected)
 }
 
+func (s *GeneratorSuite) TestInternalPackagePrologue() {
+	expected := `package mocks
+
+import baz "github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz"
+import mock "github.com/stretchr/testify/mock"
+
+`
+	generator := NewGenerator(
+		s.ctx,
+		config.Config{InPackage: false, LogLevel: "debug", ReplaceType: []string{
+			"github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz/internal/foo.InternalBaz=baz:github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz.Baz",
+		}},
+		s.getInterfaceFromFile("example_project/baz/foo.go", "Foo"),
+		pkg,
+	)
+
+	s.checkPrologueGeneration(generator, expected)
+}
+
+func (s *GeneratorSuite) TestInternalPackage() {
+	expected := `// Foo is an autogenerated mock type for the Foo type
+type Foo struct {
+	mock.Mock
+}
+
+// DoFoo provides a mock function with given fields:
+func (_m *Foo) DoFoo() string {
+	ret := _m.Called()
+
+	var r0 string
+	if rf, ok := ret.Get(0).(func() string); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(string)
+	}
+
+	return r0
+}
+
+// GetBaz provides a mock function with given fields:
+func (_m *Foo) GetBaz() (*baz.Baz, error) {
+	ret := _m.Called()
+
+	var r0 *baz.Baz
+	if rf, ok := ret.Get(0).(func() *baz.Baz); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*baz.Baz)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+type mockConstructorTestingTNewFoo interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewFoo creates a new instance of Foo. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewFoo(t mockConstructorTestingTNewFoo) *Foo {
+	mock := &Foo{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+`
+	cfg := config.Config{InPackage: false, LogLevel: "debug", ReplaceType: []string{
+		"github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz/internal/foo.InternalBaz=baz:github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz.Baz",
+	}}
+
+	s.checkGenerationWithConfig("example_project/baz/foo.go", "Foo", cfg, expected)
+}
+
 func (s *GeneratorSuite) TestGenericGenerator() {
 	expected := `// RequesterGenerics is an autogenerated mock type for the RequesterGenerics type
 type RequesterGenerics[TAny interface{}, TComparable comparable, TSigned constraints.Signed, TIntf test.GetInt, TExternalIntf io.Writer, TGenIntf test.GetGeneric[TSigned], TInlineType interface{ ~int | ~uint }, TInlineTypeGeneric interface {
@@ -2799,3 +2883,28 @@ func TestGeneratorSuite(t *testing.T) {
 	generatorSuite := new(GeneratorSuite)
 	suite.Run(t, generatorSuite)
 }
+
+func TestParseReplaceType(t *testing.T) {
+	tests := []struct {
+		value    string
+		expected replaceType
+	}{
+		{
+			value:    "github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz/internal/foo.InternalBaz",
+			expected: replaceType{alias: "", pkg: "github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz/internal/foo", typ: "InternalBaz"},
+		},
+		{
+			value:    "baz:github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz.Baz",
+			expected: replaceType{alias: "baz", pkg: "github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz", typ: "Baz"},
+		},
+		{
+			value:    "github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz",
+			expected: replaceType{alias: "", pkg: "github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz", typ: ""},
+		},
+	}
+
+	for _, test := range tests {
+		actual := parseReplaceType(test.value)
+		assert.Equal(t, test.expected, actual)
+	}
+}