From 6e967031780bc16ec83b99fee1ee705ecd3e52d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Ch=C3=A1bek?= <jakub@chabek.cz>
Date: Tue, 19 Apr 2022 19:22:52 +0200
Subject: [PATCH] Refactor mock name generation

Also fix name for unexported InPackage mocks (was Mocklowercase and now
is MockUppercase, as expected)
---
 pkg/generator.go      | 52 ++++++++++++++++++++-----------------------
 pkg/generator_test.go | 12 +++++-----
 2 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/pkg/generator.go b/pkg/generator.go
index ec0d4ed4..8da3898a 100644
--- a/pkg/generator.go
+++ b/pkg/generator.go
@@ -16,6 +16,7 @@ import (
 	"strings"
 	"text/template"
 	"unicode"
+	"unicode/utf8"
 
 	"github.com/rs/zerolog"
 	"github.com/vektra/mockery/v2/pkg/config"
@@ -195,15 +196,22 @@ func (g *Generator) getLocalizedPath(ctx context.Context, path string) string {
 	return toReturn
 }
 
-func upperFirstOnly(s string) string {
-	first := true
-	return strings.Map(func(r rune) rune {
-		if first {
-			first = false
-			return unicode.ToUpper(r)
-		}
-		return r
-	}, s)
+func (g *Generator) maybeMakeNameExported(name string, export bool) string {
+	if export && !ast.IsExported(name) {
+		return g.makeNameExported(name)
+	}
+
+	return name
+}
+
+func (g *Generator) makeNameExported(name string) string {
+	r, n := utf8.DecodeRuneInString(name)
+
+	if unicode.IsUpper(r) {
+		return name
+	}
+
+	return string(unicode.ToUpper(r)) + name[n:]
 }
 
 func (g *Generator) mockName() string {
@@ -212,17 +220,13 @@ func (g *Generator) mockName() string {
 	}
 
 	if !g.KeepTree && g.InPackage {
-		if g.Exported || ast.IsExported(g.iface.Name) {
-			return "Mock" + g.iface.Name
-		}
-
-		return "mock" + upperFirstOnly(g.iface.Name)
-	}
-	if g.Exported || ast.IsExported(g.iface.Name) {
-		return upperFirstOnly(g.iface.Name)
+		return g.maybeMakeNameExported(
+			"mock"+g.makeNameExported(g.iface.Name),
+			g.Exported || ast.IsExported(g.iface.Name),
+		)
 	}
 
-	return g.iface.Name
+	return g.maybeMakeNameExported(g.iface.Name, g.Exported)
 }
 
 func (g *Generator) expecterName() string {
@@ -715,18 +719,10 @@ func %[1]s(t testing.TB) *%[2]s {
 }
 `
 
-	funcFirstLetter := "N"
 	mockName := g.mockName()
+	constructorName := g.maybeMakeNameExported("new"+g.makeNameExported(mockName), ast.IsExported(mockName))
 
-	for _, firstLetter := range mockName {
-		if unicode.IsLower(firstLetter) {
-			funcFirstLetter = "n"
-		}
-
-		break
-	}
-
-	g.printf(constructor, funcFirstLetter+"ew"+upperFirstOnly(mockName), mockName)
+	g.printf(constructor, constructorName, mockName)
 }
 
 // generateCalled returns the Mock.Called invocation string and, if necessary, prints the
diff --git a/pkg/generator_test.go b/pkg/generator_test.go
index fd2faec9..27388be2 100644
--- a/pkg/generator_test.go
+++ b/pkg/generator_test.go
@@ -413,19 +413,19 @@ func newMockRequester_unexported(t testing.TB) *mockRequester_unexported {
 }
 
 func (s *GeneratorSuite) TestGeneratorExportedOnFlag() {
-	expected := `// Mockrequester_unexported is an autogenerated mock type for the requester_unexported type
-type Mockrequester_unexported struct {
+	expected := `// MockRequester_unexported is an autogenerated mock type for the requester_unexported type
+type MockRequester_unexported struct {
 	mock.Mock
 }
 
 // Get provides a mock function with given fields:
-func (_m *Mockrequester_unexported) Get() {
+func (_m *MockRequester_unexported) Get() {
 	_m.Called()
 }
 
-// NewMockrequester_unexported creates a new instance of Mockrequester_unexported. It also registers a cleanup function to assert the mocks expectations.
-func NewMockrequester_unexported(t testing.TB) *Mockrequester_unexported {
-	mock := &Mockrequester_unexported{}
+// NewMockRequester_unexported creates a new instance of MockRequester_unexported. It also registers a cleanup function to assert the mocks expectations.
+func NewMockRequester_unexported(t testing.TB) *MockRequester_unexported {
+	mock := &MockRequester_unexported{}
 
 	t.Cleanup(func() { mock.AssertExpectations(t) })