Support mocking functions and interfaces that use data types defined …

…in their host package

cmd/clzgenerator.go

@@ -11,7 +11,6 @@ import (
 	"strings"
 
 	"github.com/kelveny/mockcompose/pkg/gogen"
-	"github.com/kelveny/mockcompose/pkg/gosyntax"
 	"github.com/kelveny/mockcompose/pkg/logger"
 )
 
@@ -114,13 +113,15 @@ func (g *classMethodGenerator) composeMock(
 	fset *token.FileSet,
 	fnSpec *ast.FuncDecl,
 ) {
-	gosyntax.MockFunc(
+	gogen.MockFunc(
 		writer,
+		g.mockPkgName,
 		g.mockName,
 		fset,
 		fnSpec.Name.Name,
 		fnSpec.Type.Params,
 		fnSpec.Type.Results,
+		nil,
 	)
 }
 

cmd/fngenerator.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/kelveny/mockcompose/pkg/gogen"
 	"github.com/kelveny/mockcompose/pkg/gosyntax"
+	"github.com/kelveny/mockcompose/pkg/gotype"
 	"github.com/kelveny/mockcompose/pkg/logger"
 	"golang.org/x/tools/go/packages"
 )
@@ -18,6 +19,7 @@ type functionMockGenerator struct {
 	mockPkgName   string   // package name that mocking class resides
 	mockName      string   // the mocking composite class name
 	methodsToMock []string // function names that need to be mocked
+	srcPkg        string
 }
 
 // use compiler to enforce interface compliance
@@ -41,13 +43,15 @@ func (g *functionMockGenerator) generate(
 				bufWriter.Write([]byte(fmt.Sprintf("package %s\n\n", g.mockPkgName)))
 			}
 
-			gosyntax.MockFunc(
+			gogen.MockFunc(
 				bufWriter,
+				g.mockPkgName,
 				g.mockName,
 				fset,
 				fnDecl.Name.Name,
 				fnDecl.Type.Params,
 				fnDecl.Type.Results,
+				nil,
 			)
 		}
 	})
@@ -100,16 +104,24 @@ func (g *functionMockGenerator) generateViaLoadedPackage(
 				bufWriter.Write([]byte(fmt.Sprintf("package %s\n\n", g.mockPkgName)))
 
 				imports := gogen.GetPackageImports(pkg)
+				if g.srcPkg != "" {
+					imports = append(imports, gogen.ImportSpec{
+						Name: "",
+						Path: g.srcPkg,
+					})
+				}
 				gogen.WriteImportDecls(bufWriter, imports)
 			}
 
-			gosyntax.MockFunc(
+			gogen.MockFunc(
 				bufWriter,
+				g.mockPkgName,
 				g.mockName,
 				fset,
 				fnDecl.Name.Name,
 				fnDecl.Type.Params,
 				fnDecl.Type.Results,
+				gotype.FindFuncSignature(pkg, fnDecl.Name.Name),
 			)
 		}
 	})

cmd/intfgenerator.go

@@ -6,10 +6,12 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
+	"go/types"
 	"io"
 
 	"github.com/kelveny/mockcompose/pkg/gogen"
 	"github.com/kelveny/mockcompose/pkg/gosyntax"
+	"github.com/kelveny/mockcompose/pkg/gotype"
 	"github.com/kelveny/mockcompose/pkg/logger"
 	"golang.org/x/tools/go/packages"
 )
@@ -18,6 +20,7 @@ type interfaceMockGenerator struct {
 	mockPkgName string // package name that mocking class resides
 	mockName    string // the mocking composite class name
 	intfName    string // interface name
+	srcPkg      string
 }
 
 // use compiler to enforce interface compliance
@@ -33,7 +36,14 @@ func (g *interfaceMockGenerator) generate(
 		func(name string, methods []*ast.Field) {
 			if name == g.intfName {
 				imports := gogen.GetFileImports(file)
-				g.generateInterfaceMock(writer, imports, methods)
+				if g.srcPkg != "" {
+					imports = append(imports, gogen.ImportSpec{
+						Name: "",
+						Path: g.srcPkg,
+					})
+				}
+
+				g.generateInterfaceMock(writer, imports, methods, nil, name)
 			}
 		},
 	)
@@ -48,7 +58,13 @@ func (g *interfaceMockGenerator) generateViaLoadedPackage(
 		func(name string, methods []*ast.Field) {
 			if name == g.intfName {
 				imports := gogen.GetPackageImports(pkg)
-				g.generateInterfaceMock(writer, imports, methods)
+				if g.srcPkg != "" {
+					imports = append(imports, gogen.ImportSpec{
+						Name: "",
+						Path: g.srcPkg,
+					})
+				}
+				g.generateInterfaceMock(writer, imports, methods, pkg, name)
 			}
 		},
 	)
@@ -59,11 +75,13 @@ func (g *interfaceMockGenerator) generateInterfaceMock(
 	writer io.Writer,
 	imports []gogen.ImportSpec,
 	methods []*ast.Field,
+	pkg *packages.Package,
+	intfName string,
 ) error {
 	var buf bytes.Buffer
 
 	fset := token.NewFileSet()
-	if g.generateInterfaceMockInternal(&buf, fset, imports, methods) {
+	if g.generateInterfaceMockInternal(&buf, fset, imports, methods, pkg, intfName) {
 		// reload generated content to process generated code the second time
 		f, err := parser.ParseFile(fset, "", buf.Bytes(), parser.ParseComments)
 		if err != nil {
@@ -97,20 +115,33 @@ func (g *interfaceMockGenerator) generateInterfaceMockInternal(
 	fset *token.FileSet,
 	imports []gogen.ImportSpec,
 	methods []*ast.Field,
+	pkg *packages.Package,
+	intfName string,
 ) bool {
 	writer.Write([]byte(fmt.Sprintf("package %s\n\n", g.mockPkgName)))
 
 	gogen.WriteImportDecls(writer, imports)
 
 	for _, method := range methods {
 		if ftype, ok := method.Type.(*ast.FuncType); ok {
-			gosyntax.MockFunc(
+			var signature *types.Signature
+			if pkg != nil {
+				signature = gotype.FindInterfaceMethodSignature(
+					pkg,
+					intfName,
+					method.Names[0].Name,
+				)
+			}
+
+			gogen.MockFunc(
 				writer,
+				g.mockPkgName,
 				g.mockName,
 				fset,
 				method.Names[0].Name,
 				ftype.Params,
 				ftype.Results,
+				signature,
 			)
 		}
 	}

cmd/mockcompose.go

@@ -83,6 +83,7 @@ func executeOptions(options *CommandOptions) {
 			mockPkgName: options.MockPkg,
 			mockName:    options.MockName,
 			intfName:    options.IntfName,
+			srcPkg:      options.SrcPkg,
 		}
 
 		if options.SrcPkg != "" {
@@ -116,6 +117,7 @@ func executeOptions(options *CommandOptions) {
 				mockPkgName:   options.MockPkg,
 				mockName:      options.MockName,
 				methodsToMock: options.MethodsToMock,
+				srcPkg:        options.SrcPkg,
 			}
 
 			if options.SrcPkg != "" {

go.mod

@@ -4,8 +4,6 @@ go 1.16
 
 require (
 	github.com/stretchr/testify v1.7.0
-	golang.org/x/mod v0.4.2 // indirect
-	golang.org/x/sys v0.0.0-20210510120138-977fb7262007 // indirect
-	golang.org/x/tools v0.1.0
+	golang.org/x/tools v0.1.4
 	gopkg.in/yaml.v2 v2.4.0
 )

go.sum

@@ -6,30 +6,29 @@ github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210510120138-977fb7262007 h1:gG67DSER+11cZvqIMb8S8bt0vZtiN6xWYARwirrOSfE=
 golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.1.0 h1:po9/4sTYwZU9lPhi1tOrb4hCv3qrhiQ77LZfGa2OjwY=
-golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
+golang.org/x/tools v0.1.4 h1:cVngSRcfgyZCzys3KYOpCFa+4dqX/Oub9tAq00ttGVs=
+golang.org/x/tools v0.1.4/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=

pkg/gofile/helper.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/kelveny/mockcompose/pkg/logger"
+	"golang.org/x/tools/imports"
 )
 
 func DerivePackage(anchor bool) string {
@@ -91,5 +92,11 @@ func FormatGoFile(filePath string) {
 		return
 	}
 
+	bb, err = imports.Process(filePath, bb, nil)
+	if err != nil {
+		logger.Log(logger.ERROR, "Error in formatting Go imports %s, error: %s\n", filePath, err)
+		return
+	}
+
 	ioutil.WriteFile(filePath, bb, 0644)
 }