From ec95f73729c594c211c3c1a6a09916e3d6b09e2f Mon Sep 17 00:00:00 2001
From: Rangel Reale <rangelspam@gmail.com>
Date: Fri, 10 Mar 2023 09:51:33 -0300
Subject: [PATCH] use regexp to test

---
 pkg/generator_test.go | 99 +++++++++++++++++--------------------------
 1 file changed, 38 insertions(+), 61 deletions(-)

diff --git a/pkg/generator_test.go b/pkg/generator_test.go
index ea0b3456..f2d1c927 100644
--- a/pkg/generator_test.go
+++ b/pkg/generator_test.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"testing"
 
@@ -85,6 +86,37 @@ func (s *GeneratorSuite) checkGenerationWithConfig(
 	return generator
 }
 
+type regexpExpected struct {
+	shouldMatch bool
+	re          *regexp.Regexp
+}
+
+func (s *GeneratorSuite) checkGenerationRegexWithConfig(
+	filepath, interfaceName string, cfg GeneratorConfig, expected []regexpExpected,
+) *Generator {
+	generator := s.getGeneratorWithConfig(filepath, interfaceName, cfg)
+	err := generator.Generate(s.ctx)
+	s.NoError(err, "The generator ran without errors.")
+	if err != nil {
+		return generator
+	}
+	// Mirror the formatting done by normally done by golang.org/x/tools/imports in Generator.Write.
+	//
+	// While we could possibly reuse Generator.Write here in addition to Generator.Generate,
+	// it would require changing Write's signature to accept custom options, specifically to
+	// allow the fragments in preexisting cases. It's assumed that this approximation,
+	// just formatting the source, is sufficient for the needs of the current test styles.
+	var actual []byte
+	actual, fmtErr := format.Source(generator.buf.Bytes())
+	s.NoError(fmtErr, "The formatter ran without errors.")
+
+	for _, re := range expected {
+		s.Equalf(re.shouldMatch, re.re.Match(actual), "match '%s' should be %t", re.re.String(), re.shouldMatch)
+	}
+
+	return generator
+}
+
 func (s *GeneratorSuite) getGenerator(
 	filepath, interfaceName string, inPackage bool, structName string,
 ) *Generator {
@@ -2444,71 +2476,16 @@ import mock "github.com/stretchr/testify/mock"
 }
 
 func (s *GeneratorSuite) TestReplaceTypePackage() {
-	expected := `// Foo is an autogenerated mock type for the Foo type
-type Foo struct {
-	mock.Mock
-}
-
-// DoFoo provides a mock function with given fields:
-func (_m *Foo) DoFoo() string {
-	ret := _m.Called()
-
-	var r0 string
-	if rf, ok := ret.Get(0).(func() string); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Get(0).(string)
-	}
-
-	return r0
-}
-
-// GetBaz provides a mock function with given fields:
-func (_m *Foo) GetBaz() (*baz.Baz, error) {
-	ret := _m.Called()
-
-	var r0 *baz.Baz
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (*baz.Baz, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() *baz.Baz); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*baz.Baz)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-type mockConstructorTestingTNewFoo interface {
-	mock.TestingT
-	Cleanup(func())
-}
-
-// NewFoo creates a new instance of Foo. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-func NewFoo(t mockConstructorTestingTNewFoo) *Foo {
-	mock := &Foo{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}
-`
 	cfg := GeneratorConfig{InPackage: false, ReplaceType: []string{
 		"github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz/internal/foo.InternalBaz=baz:github.com/vektra/mockery/v2/pkg/fixtures/example_project/baz.Baz",
 	}}
 
-	s.checkGenerationWithConfig("example_project/baz/foo.go", "Foo", cfg, expected)
+	s.checkGenerationRegexWithConfig("example_project/baz/foo.go", "Foo", cfg, []regexpExpected{
+		// func (_m *Foo) GetBaz() (*baz.Baz, error)
+		{true, regexp.MustCompile(`func \([^\)]+\) GetBaz\(\) \(\*baz\.Baz`)},
+		// func (_m *Foo) GetBaz() (*foo.InternalBaz, error)
+		{false, regexp.MustCompile(`func \([^\)]+\) GetBaz\(\) \(\*foo\.InternalBaz`)},
+	})
 }
 
 func (s *GeneratorSuite) TestGenericGenerator() {