Merge pull request #710 from RangelReale/replace-types-specific

Fix replace-type for different packages from the same source
vektra · Oct 5, 2023 · 246df9d · 246df9d
2 parents f54eea9 + ee33b3d
commit 246df9d
Show file tree

Hide file tree

Showing 7 changed files with 91 additions and 9 deletions.
diff --git a/pkg/fixtures/example_project/replace_type/README.md b/pkg/fixtures/example_project/replace_type/README.md
@@ -0,0 +1,9 @@
+## Fix replace-type for different packages from the same source
+
+[Issue 710](https://github.com/vektra/mockery/pull/710)
+
+This package is used to test the case where multiple types come from the same package (`replace_type/rti/internal`),
+but results in types in different packages (`replace_type/rt1` and `replace_type/rt2`).
+
+Tests `TestReplaceTypePackageMultiplePrologue` and `TestReplaceTypePackageMultiple` use it to check if this outputs
+the correct import and type names.
diff --git a/pkg/fixtures/example_project/replace_type/rt.go b/pkg/fixtures/example_project/replace_type/rt.go
@@ -0,0 +1,11 @@
+package replace_type
+
+import (
+	"github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt1"
+	"github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt2"
+)
+
+type RType interface {
+	Replace1(f rt1.RType1)
+	Replace2(f rt2.RType2)
+}
diff --git a/pkg/fixtures/example_project/replace_type/rti/internal/rti.go b/pkg/fixtures/example_project/replace_type/rti/internal/rti.go
@@ -0,0 +1,9 @@
+package internal
+
+type RTInternal1 struct {
+	A int
+}
+
+type RTInternal2 struct {
+	B string
+}
diff --git a/pkg/fixtures/example_project/replace_type/rti/rt1/rt1.go b/pkg/fixtures/example_project/replace_type/rti/rt1/rt1.go
@@ -0,0 +1,5 @@
+package rt1
+
+import "github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/internal"
+
+type RType1 = internal.RTInternal1
diff --git a/pkg/fixtures/example_project/replace_type/rti/rt2/rt2.go b/pkg/fixtures/example_project/replace_type/rti/rt2/rt2.go
@@ -0,0 +1,5 @@
+package rt2
+
+import "github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/internal"
+
+type RType2 = internal.RTInternal2
diff --git a/pkg/generator.go b/pkg/generator.go
@@ -103,7 +103,7 @@ func NewGenerator(ctx context.Context, c GeneratorConfig, iface *Interface, pkg
 	}
 
 	g.parseReplaceTypes(ctx)
-	g.addPackageImportWithName(ctx, "github.com/stretchr/testify/mock", "mock")
+	g.addPackageImportWithName(ctx, "github.com/stretchr/testify/mock", "mock", nil)
 
 	return g
 }
@@ -161,7 +161,7 @@ func (g *Generator) getPackageScopedType(ctx context.Context, o *types.TypeName)
 		(!g.config.KeepTree && g.config.InPackage && o.Pkg() == g.iface.Pkg) {
 		return o.Name()
 	}
-	pkg := g.addPackageImport(ctx, o.Pkg())
+	pkg := g.addPackageImport(ctx, o.Pkg(), o)
 	name := o.Name()
 	g.checkReplaceType(ctx, func(from *replaceType, to *replaceType) bool {
 		if o.Pkg().Path() == from.pkg && name == from.typ {
@@ -173,23 +173,28 @@ func (g *Generator) getPackageScopedType(ctx context.Context, o *types.TypeName)
 	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) addPackageImport(ctx context.Context, pkg *types.Package, o *types.TypeName) string {
+	return g.addPackageImportWithName(ctx, pkg.Path(), pkg.Name(), o)
 }
 
 func (g *Generator) checkReplaceType(ctx context.Context, f func(from *replaceType, to *replaceType) bool) {
-	for _, item := range g.replaceTypeCache {
-		if !f(item.from, item.to) {
-			break
+	// check most specific first
+	for _, hasType := range []bool{true, false} {
+		for _, item := range g.replaceTypeCache {
+			if (item.from.typ != "") == hasType {
+				if !f(item.from, item.to) {
+					break
+				}
+			}
 		}
 	}
 }
 
-func (g *Generator) addPackageImportWithName(ctx context.Context, path, name string) string {
+func (g *Generator) addPackageImportWithName(ctx context.Context, path, name string, o *types.TypeName) string {
 	log := zerolog.Ctx(ctx)
 	replaced := false
 	g.checkReplaceType(ctx, func(from *replaceType, to *replaceType) bool {
-		if path == from.pkg {
+		if o != nil && path == from.pkg && (from.typ == "" || o.Name() == from.typ) {
 			log.Debug().Str("from", path).Str("to", to.pkg).Msg("changing package path")
 			replaced = true
 			path = to.pkg

diff --git a/pkg/generator_test.go b/pkg/generator_test.go
@@ -711,6 +711,44 @@ func (s *GeneratorSuite) TestReplaceTypePackage() {
 	})
 }
 
+func (s *GeneratorSuite) TestReplaceTypePackageMultiplePrologue() {
+	expected := `package mocks
+
+import mock "github.com/stretchr/testify/mock"
+import replace_type "github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type"
+import rt1 "github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt1"
+import rt2 "github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt2"
+
+`
+	generator := NewGenerator(
+		s.ctx,
+		GeneratorConfig{InPackage: false, ReplaceType: []string{
+			"github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/internal.RTInternal1=rt1:github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt1.RType1",
+			"github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/internal.RTInternal2=rt2:github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt2.RType2",
+		}},
+		s.getInterfaceFromFile("example_project/replace_type/rt.go", "RType"),
+		pkg,
+	)
+
+	s.checkPrologueGeneration(generator, expected)
+}
+
+func (s *GeneratorSuite) TestReplaceTypePackageMultiple() {
+	cfg := GeneratorConfig{InPackage: false, ReplaceType: []string{
+		// first one will be ignored by the more specific ones
+		"github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/internal=fiz3:github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt3",
+		"github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/internal.RTInternal1=rt1:github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt1.RType1",
+		"github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/internal.RTInternal2=rt2:github.com/vektra/mockery/v2/pkg/fixtures/example_project/replace_type/rti/rt2.RType2",
+	}}
+
+	s.checkGenerationRegexWithConfig("example_project/replace_type/rt.go", "RType", cfg, []regexpExpected{
+		// func (_m *RType) Replace1(f rt1.RType1)
+		{true, regexp.MustCompile(`func \([^\)]+\) Replace1\(f rt1\.RType1`)},
+		// func (_m *RType) Replace2(f rt2.RType2)
+		{true, regexp.MustCompile(`func \([^\)]+\) Replace2\(f rt2\.RType2`)},
+	})
+}
+
 func (s *GeneratorSuite) TestGenericGenerator() {
 	s.checkGeneration("generic.go", "RequesterGenerics", false, "", "")
 }