Add disable-func-mocks parameter

Helps with #716
vektra · Aug 20, 2024 · 56379c8 · 56379c8
1 parent c6d2a9e
commit 56379c8
Show file tree

Hide file tree

Showing 7 changed files with 75 additions and 20 deletions.
diff --git a/cmd/mockery.go b/cmd/mockery.go
@@ -83,6 +83,7 @@ func NewRootCmd() *cobra.Command {
 	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")
+	pFlags.Bool("disable-func-mocks", false, "Disable generation of function mocks.")
 
 	if err := viperCfg.BindPFlags(pFlags); err != nil {
 		panic(fmt.Sprintf("failed to bind PFlags: %v", err))
@@ -238,7 +239,7 @@ func (r *RootApp) Run() error {
 		if err != nil {
 			return fmt.Errorf("failed to get package from config: %w", err)
 		}
-		parser := pkg.NewParser(buildTags)
+		parser := pkg.NewParser(buildTags, pkg.ParserDisableFuncMocks(r.Config.DisableFuncMocks))
 
 		if err := parser.ParsePackages(ctx, configuredPackages); err != nil {
 			log.Error().Err(err).Msg("unable to parse packages")

diff --git a/docs/configuration.md b/docs/configuration.md
@@ -64,6 +64,7 @@ Parameter Descriptions
 | `config`                                               | :fontawesome-solid-x:     | `#!yaml ""`                           | Set the location of the mockery config file.                                                                                                                                                                                                         |
 | `dir`                                                  | :fontawesome-solid-check: | `#!yaml "mocks/{{.PackagePath}}"`     | The directory where the mock file will be outputted to.                                                                                                                                                                                              |
 | `disable-config-search`                                | :fontawesome-solid-x:     | `#!yaml false`                        | Disable searching for configuration files                                                                                                                                                                                                            |
+| `disable-func-mocks`                                   | :fontawesome-solid-x:     | `#!yaml false`                        | Disable generation of function mocks.                                                                                                                                                                                                                |
 | `disable-version-string`                               | :fontawesome-solid-x:     | `#!yaml false`                        | Disable the version string in the generated mock files.                                                                                                                                                                                              |
 | `dry-run`                                              | :fontawesome-solid-x:     | `#!yaml false`                        | Print the actions that would be taken, but don't perform the actions.                                                                                                                                                                                |
 | `exclude`                                              | :fontawesome-solid-x:     | `#!yaml []`                           | Specify subpackages to exclude when using `#!yaml recursive: True`                                                                                                                                                                                   |

diff --git a/e2e/.mockery-disable-func-mock.yaml b/e2e/.mockery-disable-func-mock.yaml
@@ -0,0 +1,6 @@
+all: false
+log-level: info
+packages:
+  github.com/vektra/mockery/v2/pkg/fixtures:
+    interfaces:
+      SendFunc:
diff --git a/e2e/run_all.sh b/e2e/run_all.sh
@@ -1,4 +1,5 @@
 #!/bin/bash
+set -e
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 
 for file in $(ls $SCRIPT_DIR/test_*.sh); do

diff --git a/e2e/test_disable_func_mocks.sh b/e2e/test_disable_func_mocks.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+export MOCKERY_CONFIG="e2e/.mockery-disable-func-mock.yaml"
+export MOCKERY_LOG_LEVEL="error"
+
+MOCKERY_DISABLE_FUNC_MOCKS="false" go run github.com/go-task/task/v3/cmd/task mocks.generate
+
+if [ -f "./mocks/github.com/vektra/mockery/v2/pkg/fixtures/mock_SendFunc.go" ]; then
+    echo "file exists as expected"
+else
+    echo "file doesn't exist when we expected it to exist"
+    exit 1
+fi
+
+go run github.com/go-task/task/v3/cmd/task mocks.remove
+MOCKERY_DISABLE_FUNC_MOCKS="true" go run github.com/go-task/task/v3/cmd/task mocks.generate
+if [ -f "./mocks/github.com/vektra/mockery/v2/pkg/fixtures/mock_SendFunc.go" ]; then
+    echo "SendFunc mock exists when we expected it to not be generated."
+    exit 1
+else
+    echo "SendFunc mock doesn't exist as expected"
+fi
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -41,6 +41,7 @@ type Config struct {
 	Cpuprofile           string                 `mapstructure:"cpuprofile"`
 	Dir                  string                 `mapstructure:"dir"`
 	DisableConfigSearch  bool                   `mapstructure:"disable-config-search"`
+	DisableFuncMocks     bool                   `mapstructure:"disable-func-mocks"`
 	DisableVersionString bool                   `mapstructure:"disable-version-string"`
 	DryRun               bool                   `mapstructure:"dry-run"`
 	ExcludeRegex         string                 `mapstructure:"exclude-regex"`

diff --git a/pkg/parse.go b/pkg/parse.go
@@ -17,14 +17,15 @@ import (
 )
 
 type fileEntry struct {
-	fileName   string
-	pkg        *packages.Package
-	syntax     *ast.File
-	interfaces []string
+	fileName         string
+	pkg              *packages.Package
+	syntax           *ast.File
+	interfaces       []string
+	disableFuncMocks bool
 }
 
 func (f *fileEntry) ParseInterfaces(ctx context.Context) {
-	nv := NewNodeVisitor(ctx)
+	nv := NewNodeVisitor(ctx, f.disableFuncMocks)
 	ast.Walk(nv, f.syntax)
 	f.interfaces = nv.DeclaredInterfaces()
 }
@@ -40,9 +41,15 @@ type Parser struct {
 	parserPackages    []*types.Package
 	conf              packages.Config
 	packageLoadCache  map[string]packageLoadEntry
+	disableFuncMocks  bool
 }
 
-func NewParser(buildTags []string) *Parser {
+func ParserDisableFuncMocks(disable bool) func(*Parser) {
+	return func(p *Parser) {
+		p.disableFuncMocks = disable
+	}
+}
+func NewParser(buildTags []string, opts ...func(*Parser)) *Parser {
 	var conf packages.Config
 	conf.Mode = packages.NeedTypes |
 		packages.NeedTypesSizes |
@@ -56,12 +63,16 @@ func NewParser(buildTags []string) *Parser {
 	if len(buildTags) > 0 {
 		conf.BuildFlags = []string{"-tags", strings.Join(buildTags, ",")}
 	}
-	return &Parser{
+	p := &Parser{
 		parserPackages:    make([]*types.Package, 0),
 		entriesByFileName: map[string]*fileEntry{},
 		conf:              conf,
 		packageLoadCache:  map[string]packageLoadEntry{},
 	}
+	for _, opt := range opts {
+		opt(p)
+	}
+	return p
 }
 
 func (p *Parser) loadPackages(fpath string) ([]*packages.Package, error) {
@@ -93,9 +104,10 @@ func (p *Parser) ParsePackages(ctx context.Context, packageNames []string) error
 				Str("file", file).
 				Msgf("found file")
 			entry := fileEntry{
-				fileName: file,
-				pkg:      pkg,
-				syntax:   pkg.Syntax[fileIdx],
+				fileName:         file,
+				pkg:              pkg,
+				syntax:           pkg.Syntax[fileIdx],
+				disableFuncMocks: p.disableFuncMocks,
 			}
 			entry.ParseInterfaces(ctx)
 			p.files = append(p.files, &entry)
@@ -321,14 +333,15 @@ func (s sortableIFaceList) Less(i, j int) bool {
 }
 
 type NodeVisitor struct {
-	declaredInterfaces            []string
-	genericInstantiationInterface map[string]any
-	ctx                           context.Context
+	declaredInterfaces []string
+	disableFuncMocks   bool
+	ctx                context.Context
 }
 
-func NewNodeVisitor(ctx context.Context) *NodeVisitor {
+func NewNodeVisitor(ctx context.Context, disableFuncMocks bool) *NodeVisitor {
 	return &NodeVisitor{
 		declaredInterfaces: make([]string, 0),
+		disableFuncMocks:   disableFuncMocks,
 		ctx:                ctx,
 	}
 }
@@ -337,6 +350,14 @@ func (nv *NodeVisitor) DeclaredInterfaces() []string {
 	return nv.declaredInterfaces
 }
 
+func (nv *NodeVisitor) add(ctx context.Context, n *ast.TypeSpec) {
+	log := zerolog.Ctx(ctx)
+	log.Debug().
+		Str("node-type", fmt.Sprintf("%T", n.Type)).
+		Msg("found node with acceptable type for mocking")
+	nv.declaredInterfaces = append(nv.declaredInterfaces, n.Name.Name)
+}
+
 func (nv *NodeVisitor) Visit(node ast.Node) ast.Visitor {
 	log := zerolog.Ctx(nv.ctx)
 
@@ -348,11 +369,13 @@ func (nv *NodeVisitor) Visit(node ast.Node) ast.Visitor {
 			Logger()
 
 		switch n.Type.(type) {
-		case *ast.InterfaceType, *ast.FuncType, *ast.IndexExpr:
-			log.Debug().
-				Str("node-type", fmt.Sprintf("%T", n.Type)).
-				Msg("found node with acceptable type for mocking")
-			nv.declaredInterfaces = append(nv.declaredInterfaces, n.Name.Name)
+		case *ast.FuncType:
+			if nv.disableFuncMocks {
+				break
+			}
+			nv.add(nv.ctx, n)
+		case *ast.InterfaceType, *ast.IndexExpr:
+			nv.add(nv.ctx, n)
 		default:
 			log.Debug().Msg("Found node with unacceptable type for mocking. Rejecting.")
 		}