Fine tune and README.md update

cmd/clzgenerator.go

@@ -14,6 +14,8 @@ import (
 
 	"github.com/kelveny/mockcompose/pkg/gogen"
 	"github.com/kelveny/mockcompose/pkg/gosyntax"
+	gosyntaxtyp "github.com/kelveny/mockcompose/pkg/gosyntax"
+
 	"github.com/kelveny/mockcompose/pkg/gotype"
 	"github.com/kelveny/mockcompose/pkg/logger"
 )
@@ -35,6 +37,46 @@ type classMethodGenerator struct {
 	methodsToMock  []string // method function names that need to be mocked
 }
 
+type generatorContext struct {
+	mockedFunctions map[string]any
+
+	// class type declaraion string -> method name -> *ReceiverSpec
+	clzMethods map[string]map[string]*gosyntax.ReceiverSpec
+}
+
+func (c *generatorContext) hasFunctionMocked(fnName string) bool {
+	if len(c.mockedFunctions) > 0 {
+		if _, ok := c.mockedFunctions[fnName]; ok {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *generatorContext) recordMockedFunction(fnName string) {
+	if c.mockedFunctions == nil {
+		c.mockedFunctions = make(map[string]any)
+	}
+	c.mockedFunctions[fnName] = struct{}{}
+}
+
+//go:generate mockcompose -n gctx_findClassMethods -c generatorContext -real findClassMethods,gosyntax
+func (c *generatorContext) findClassMethods(
+	clzTypeDeclString string,
+	fset *token.FileSet,
+	f *ast.File,
+) map[string]*gosyntaxtyp.ReceiverSpec {
+	if c.clzMethods == nil {
+		c.clzMethods = make(map[string]map[string]*gosyntaxtyp.ReceiverSpec)
+	}
+
+	if _, ok := c.clzMethods[clzTypeDeclString]; !ok {
+		c.clzMethods[clzTypeDeclString] = gosyntax.FindClassMethods(clzTypeDeclString, fset, f)
+	}
+
+	return c.clzMethods[clzTypeDeclString]
+}
+
 // use compiler to enforce interface compliance
 var _ parsedFileGenerator = (*classMethodGenerator)(nil)
 
@@ -195,12 +237,12 @@ func (g *classMethodGenerator) getAutoMockCalleeConfig(
 }
 
 func (g *classMethodGenerator) composeMock(
-	composeContext map[string]any,
+	generatorCtx *generatorContext,
 	writer io.Writer,
 	fset *token.FileSet,
 	fnSpec *ast.FuncDecl,
 ) {
-	if _, ok := composeContext[fnSpec.Name.Name]; !ok {
+	if !generatorCtx.hasFunctionMocked(fnSpec.Name.Name) {
 		gogen.MockFunc(
 			writer,
 			g.mockPkgName,
@@ -211,7 +253,8 @@ func (g *classMethodGenerator) composeMock(
 			fnSpec.Type.Results,
 			nil,
 		)
-		composeContext[fnSpec.Name.Name] = struct{}{}
+
+		generatorCtx.recordMockedFunction(fnSpec.Name.Name)
 	}
 }
 
@@ -266,7 +309,7 @@ func (g *classMethodGenerator) generateInternal(
 ) (generated bool, autoMockPkgs []string) {
 	writer.Write([]byte(fmt.Sprintf("package %s\n\n", g.mockPkgName)))
 
-	composeContext := map[string]any{}
+	generatorCtx := &generatorContext{}
 	imports := gosyntax.GetFileImportsAsMap(file)
 
 	if len(file.Decls) > 0 {
@@ -286,7 +329,7 @@ func (g *classMethodGenerator) generateInternal(
 						//
 
 						// find out callee situation
-						clzMethods := gosyntax.FindClassMethods(receiverSpec.TypeDecl, fset, file)
+						clzMethods := generatorCtx.findClassMethods(receiverSpec.TypeDecl, fset, file)
 						v := gosyntax.NewCalleeVisitor(
 							imports,
 							clzMethods,
@@ -314,7 +357,7 @@ func (g *classMethodGenerator) generateInternal(
 						autoMockPeer, pkgs := g.getAutoMockCalleeConfig(fnSpec.Name.Name)
 						if autoMockPeer {
 							// peer callee in order of how it is declared in file
-							g.generateMethodPeerCallees(composeContext, writer, fset, file, fnSpec, v)
+							g.generateMethodPeerCallees(generatorCtx, writer, fset, file, fnSpec, v)
 						}
 
 						if len(pkgs) > 0 {
@@ -362,7 +405,7 @@ func (g *classMethodGenerator) generateInternal(
 					}
 				} else if matchType == MATCH_MOCK {
 					// generate mocked method
-					g.composeMock(composeContext, writer, fset, fnSpec)
+					g.composeMock(generatorCtx, writer, fset, fnSpec)
 				}
 			} else {
 				// for any non-function declaration, export only imports
@@ -378,7 +421,7 @@ func (g *classMethodGenerator) generateInternal(
 }
 
 func (g *classMethodGenerator) generateMethodPeerCallees(
-	composeContext map[string]any,
+	generatorCtx *generatorContext,
 	writer io.Writer,
 	fset *token.FileSet,
 	file *ast.File,
@@ -394,7 +437,7 @@ func (g *classMethodGenerator) generateMethodPeerCallees(
 				gosyntax.ForEachFuncDeclInFile(file, func(fnSpec *ast.FuncDecl) {
 					if fnSpec.Name.Name == peerMethod &&
 						gosyntax.ReceiverDeclString(fset, callerFnSpec.Recv) == gosyntax.ReceiverDeclString(fset, fnSpec.Recv) {
-						g.composeMock(composeContext, writer, fset, fnSpec)
+						g.composeMock(generatorCtx, writer, fset, fnSpec)
 					}
 				})
 			}

cmd/clzgenerator_test.go

@@ -0,0 +1,80 @@
+package cmd
+
+import (
+	"testing"
+
+	"github.com/kelveny/mockcompose/pkg/gosyntax"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_generatorContext_findClassMethods_caching(t *testing.T) {
+	assert := require.New(t)
+
+	g := &gctx_findClassMethods{}
+
+	g.mock_gctx_findClassMethods_findClassMethods_gosyntax.On(
+		"FindClassMethods",
+		mock.Anything,
+		mock.Anything,
+		mock.Anything,
+	).Return(
+		map[string]*gosyntax.ReceiverSpec{
+			"Foo": {
+				Name:     "f",
+				TypeDecl: "*foo",
+			},
+		},
+	)
+
+	// call it once
+	methods := g.findClassMethods("*foo", nil, nil)
+	assert.EqualValues(
+		map[string]*gosyntax.ReceiverSpec{
+			"Foo": {
+				Name:     "f",
+				TypeDecl: "*foo",
+			},
+		},
+		methods,
+	)
+
+	// call it the second time
+	methods = g.findClassMethods("*foo", nil, nil)
+	assert.EqualValues(
+		map[string]*gosyntax.ReceiverSpec{
+			"Foo": {
+				Name:     "f",
+				TypeDecl: "*foo",
+			},
+		},
+		methods,
+	)
+
+	// assert on caching behave
+	g.mock_gctx_findClassMethods_findClassMethods_gosyntax.AssertNumberOfCalls(t, "FindClassMethods", 1)
+}
+
+func Test_generatorContext_findClassMethods_nil_return(t *testing.T) {
+	assert := require.New(t)
+
+	g := &gctx_findClassMethods{}
+
+	g.mock_gctx_findClassMethods_findClassMethods_gosyntax.On(
+		"FindClassMethods",
+		mock.Anything,
+		mock.Anything,
+		mock.Anything,
+	).Return(nil)
+
+	// call it once
+	methods := g.findClassMethods("*foo", nil, nil)
+	assert.Nil(methods)
+
+	// call it the second time
+	methods = g.findClassMethods("*foo", nil, nil)
+	assert.Nil(methods)
+
+	// assert on caching behave
+	g.mock_gctx_findClassMethods_findClassMethods_gosyntax.AssertNumberOfCalls(t, "FindClassMethods", 1)
+}

cmd/mockc_gctx_findClassMethods_test.go

@@ -0,0 +1,52 @@
+// CODE GENERATED AUTOMATICALLY WITH github.com/kelveny/mockcompose
+// THIS FILE SHOULD NOT BE EDITED BY HAND
+package cmd
+
+import (
+	"go/ast"
+	"go/token"
+
+	"github.com/kelveny/mockcompose/pkg/gosyntax"
+	gosyntaxtyp "github.com/kelveny/mockcompose/pkg/gosyntax"
+	"github.com/stretchr/testify/mock"
+)
+
+type gctx_findClassMethods struct {
+	generatorContext
+	mock.Mock
+	mock_gctx_findClassMethods_findClassMethods_gosyntax
+}
+
+type mock_gctx_findClassMethods_findClassMethods_gosyntax struct {
+	mock.Mock
+}
+
+func (c *gctx_findClassMethods) findClassMethods(clzTypeDeclString string, fset *token.FileSet, f *ast.File) map[string]*gosyntaxtyp.ReceiverSpec {
+	gosyntax := &c.mock_gctx_findClassMethods_findClassMethods_gosyntax
+
+	if c.clzMethods == nil {
+		c.clzMethods = make(map[string]map[string]*gosyntaxtyp.ReceiverSpec)
+	}
+	if _, ok := c.clzMethods[clzTypeDeclString]; !ok {
+		c.clzMethods[clzTypeDeclString] = gosyntax.FindClassMethods(clzTypeDeclString, fset, f)
+	}
+	return c.clzMethods[clzTypeDeclString]
+}
+
+func (m *mock_gctx_findClassMethods_findClassMethods_gosyntax) FindClassMethods(clzTypeDeclString string, fset *token.FileSet, f *ast.File) map[string]*gosyntax.ReceiverSpec {
+
+	_mc_ret := m.Called(clzTypeDeclString, fset, f)
+
+	var _r0 map[string]*gosyntax.ReceiverSpec
+
+	if _rfn, ok := _mc_ret.Get(0).(func(string, *token.FileSet, *ast.File) map[string]*gosyntax.ReceiverSpec); ok {
+		_r0 = _rfn(clzTypeDeclString, fset, f)
+	} else {
+		if _mc_ret.Get(0) != nil {
+			_r0 = _mc_ret.Get(0).(map[string]*gosyntax.ReceiverSpec)
+		}
+	}
+
+	return _r0
+
+}

pkg/gosyntax/helper.go

@@ -429,16 +429,16 @@ type ReceiverSpec struct {
 
 // Find all methods of a "class"
 //
-// For methods with pointer receivers, prepend "*" to type name when passed in clzName
-func FindClassMethods(clzName string, fset *token.FileSet, f *ast.File) map[string]*ReceiverSpec {
+// For methods with pointer receivers, prepend "*" to type name when passed in clzTypeDeclString
+func FindClassMethods(clzTypeDeclString string, fset *token.FileSet, f *ast.File) map[string]*ReceiverSpec {
 	methods := make(map[string]*ReceiverSpec)
 
 	ForEachFuncDeclInFile(f, func(funcDecl *ast.FuncDecl) {
 		if funcDecl.Recv != nil {
 			recvStr := ParamListDeclString(fset, funcDecl.Recv)
 			tokens := strings.Split(recvStr, " ")
 			if len(tokens) == 2 {
-				if tokens[1] == clzName {
+				if tokens[1] == clzTypeDeclString {
 					methods[funcDecl.Name.Name] = &ReceiverSpec{
 						Name:     tokens[0],
 						TypeDecl: tokens[1],

test/mix/mix_receiver.go

@@ -0,0 +1,39 @@
+package mix
+
+type mixReceiver struct {
+	value string
+}
+
+func (v mixReceiver) getValue() string {
+	return v.value
+}
+
+func (p *mixReceiver) setValue(val string) {
+	p.value = val
+}
+
+// only callees with the same receiver type (either by-value or by-reference type)
+// will be considered as peer callee method
+//
+// in checkAndSet method, since receiver type of checkAndSet is by-reference,
+// only by-reference setValue will be considered as its peer callee method automatically.
+//
+// in checkAndSetOnTarget, since receiver type of checkAndSetOnTarget is by-value
+// only by-value getValue will be considered as its peer callee method automatically.
+
+//go:generate mockcompose -n mix_checkAndSet -c mixReceiver -real checkAndSet,this
+func (p *mixReceiver) checkAndSet(s string, val string) {
+	if p.getValue() != s {
+		p.setValue(val)
+	}
+}
+
+//go:generate mockcompose -n mix_checkAndSetOnTarget -c mixReceiver -real checkAndSetOnTarget,this
+func (v mixReceiver) checkAndSetOnTarget(p *mixReceiver, s string, val string) {
+	if v.getValue() != s {
+
+		// use it for code generation validation, has no effect on v
+		v.setValue(val)
+		p.setValue(val)
+	}
+}

test/mix/mockc_mix_checkAndSetOnTarget_test.go

@@ -0,0 +1,37 @@
+// CODE GENERATED AUTOMATICALLY WITH github.com/kelveny/mockcompose
+// THIS FILE SHOULD NOT BE EDITED BY HAND
+package mix
+
+import (
+	"github.com/stretchr/testify/mock"
+)
+
+type mix_checkAndSetOnTarget struct {
+	mixReceiver
+	mock.Mock
+}
+
+func (v mix_checkAndSetOnTarget) checkAndSetOnTarget(p *mixReceiver, s string, val string) {
+	if v.getValue() != s {
+		v.setValue(val)
+		p.setValue(val)
+	}
+}
+
+func (m *mix_checkAndSetOnTarget) getValue() string {
+
+	_mc_ret := m.Called()
+
+	var _r0 string
+
+	if _rfn, ok := _mc_ret.Get(0).(func() string); ok {
+		_r0 = _rfn()
+	} else {
+		if _mc_ret.Get(0) != nil {
+			_r0 = _mc_ret.Get(0).(string)
+		}
+	}
+
+	return _r0
+
+}

test/mix/mockc_mix_checkAndSet_test.go

@@ -0,0 +1,24 @@
+// CODE GENERATED AUTOMATICALLY WITH github.com/kelveny/mockcompose
+// THIS FILE SHOULD NOT BE EDITED BY HAND
+package mix
+
+import (
+	"github.com/stretchr/testify/mock"
+)
+
+type mix_checkAndSet struct {
+	mixReceiver
+	mock.Mock
+}
+
+func (p *mix_checkAndSet) checkAndSet(s string, val string) {
+	if p.getValue() != s {
+		p.setValue(val)
+	}
+}
+
+func (m *mix_checkAndSet) setValue(val string) {
+
+	m.Called(val)
+
+}