From 7ba373f428567bfec972c2b3ce32c642beb2c314 Mon Sep 17 00:00:00 2001
From: Ross Hammermeister <glitch17@gmail.com>
Date: Mon, 22 Jun 2020 14:47:54 -0600
Subject: [PATCH] Fixing #291 where a different import with the same name as
 the package would not get included in the mock

---
 pkg/generator.go      |  7 +------
 pkg/generator_test.go | 37 +++++++++++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/pkg/generator.go b/pkg/generator.go
index cd48a621..84aebaac 100644
--- a/pkg/generator.go
+++ b/pkg/generator.go
@@ -259,17 +259,12 @@ func (g *Generator) generateImports(ctx context.Context) {
 	log.Debug().Msgf("generating imports")
 	log.Debug().Msgf("%v", g.nameToPackagePath)
 
-	pkgPath := g.nameToPackagePath[g.iface.Pkg.Name()]
 	// Sort by import name so that we get a deterministic order
 	for _, name := range g.sortedImportNames() {
 		logImport := log.With().Str(logging.LogKeyImport, g.nameToPackagePath[name]).Logger()
 		logImport.Debug().Msgf("found import")
 
 		path := g.nameToPackagePath[name]
-		if g.InPackage && path == pkgPath {
-			logImport.Debug().Msgf("import (%s) equals interface's package path (%s), skipping", path, pkgPath)
-			continue
-		}
 		g.printf("import %s \"%s\"\n", name, path)
 	}
 }
@@ -659,7 +654,7 @@ func (g *Generator) generateCalled(list *paramList, formattedParamNames string)
 }
 
 func (g *Generator) Write(w io.Writer) error {
-	opt := &imports.Options{Comments: true}
+	opt := &imports.Options{Comments: true, FormatOnly: true}
 	theBytes := g.buf.Bytes()
 
 	res, err := imports.Process("mock.go", theBytes, opt)
diff --git a/pkg/generator_test.go b/pkg/generator_test.go
index 884764a0..80c269e6 100644
--- a/pkg/generator_test.go
+++ b/pkg/generator_test.go
@@ -14,7 +14,7 @@ import (
 	"github.com/vektra/mockery/pkg/config"
 )
 
-const pkg = "test"
+const pkg = "pkg/test"
 
 type GeneratorSuite struct {
 	suite.Suite
@@ -962,7 +962,7 @@ func (s *GeneratorSuite) TestGeneratorWithImportSameAsLocalPackageInpkgNoCycle()
 		InPackage: true,
 	}, iface, pkg)
 	gen.GeneratePrologue(s.ctx, pkg)
-	s.NotContains(gen.buf.String(), `import test "github.com/vektra/mockery/pkg/fixtures/test"`)
+	s.Contains(gen.buf.String(), `import test "github.com/vektra/mockery/pkg/fixtures/test"`)
 }
 
 func (s *GeneratorSuite) TestMapToInterface() {
@@ -1099,6 +1099,39 @@ func (_m *A) Call() (test.B, error) {
 	)
 }
 
+func (s *GeneratorSuite) TestGeneratorForStructValueReturnInPackage() {
+	expected := `// MockA is an autogenerated mock type for the A type
+type MockA struct {
+	mock.Mock
+}
+
+// Call provides a mock function with given fields:
+func (_m *MockA) Call() (B, error) {
+	ret := _m.Called()
+
+	var r0 B
+	if rf, ok := ret.Get(0).(func() B); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(B)
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+`
+	s.checkGeneration(
+		filepath.Join(fixturePath, "struct_value.go"), "A", true, "",
+		expected,
+	)
+}
+
 func (s *GeneratorSuite) TestStructNameOverride() {
 	expected := `// Requester2OverrideName is an autogenerated mock type for the Requester2 type
 type Requester2OverrideName struct {