Add support for embedding types in mocks

README.md

@@ -46,8 +46,15 @@ moq [flags] source-dir interface [interface2 [interface3 [...]]]
     	generate functions to facilitate resetting calls made to a mock
 
 Specifying an alias for the mock is also supported with the format 'interface:alias'
-
 Ex: moq -pkg different . MyInterface:MyMock
+
+To embed types into the mock, use the format 'interface{type1,type2...}' or 'interface:alias{type1,type2...}'
+- The types must be declared in source package
+- You can embed a pointer type by using the '*' prefix
+Ex1: moq -pkg different src MyInterface{Type1}
+Ex2: moq -pkg different src MyInterface{*Type1}
+Ex3: moq -pkg different src MyInterface{Type1,Type2}
+Ex4: moq -pkg different src MyInterface:MyMock{Type1}
 ```
 
 **NOTE:** `source-dir` is the directory where the source code (definition) of the target interface is located.

internal/template/template.go

@@ -89,6 +89,13 @@ type {{.MockName}}
 			{{- if $index}}, {{end}}{{$param.Name | Exported}} {{$param.TypeString}}
 	{{- end -}}]
 {{- end }} struct {
+{{- range .EmbeddedTypes}}
+	{{if .IsPointer}}*{{end -}}
+	{{$.SrcPkgQualifier -}}
+	{{- .Name -}}
+{{end}}
+{{- if .EmbeddedTypes}}
+{{end -}}
 {{- range .Methods}}
 	// {{.Name}}Func mocks the {{.Name}} method.
 	{{.Name}}Func func({{.ArgList}}) {{.ReturnArgTypeList}}

internal/template/template_data.go

@@ -37,6 +37,7 @@ type MockData struct {
 	MockName      string
 	TypeParams    []TypeParamData
 	Methods       []MethodData
+	EmbeddedTypes []EmbeddedTypeData
 }
 
 // MethodData is the data which represents a method on some interface.
@@ -132,3 +133,10 @@ func (p ParamData) CallName() string {
 func (p ParamData) TypeString() string {
 	return p.Var.TypeString()
 }
+
+// EmbeddedTypeData is the data which represents a type will be embedded
+// in the mock.
+type EmbeddedTypeData struct {
+	Name      string
+	IsPointer bool
+}

main.go

@@ -43,8 +43,17 @@ func main() {
 	flag.Usage = func() {
 		fmt.Println(`moq [flags] source-dir interface [interface2 [interface3 [...]]]`)
 		flag.PrintDefaults()
+		fmt.Println()
 		fmt.Println(`Specifying an alias for the mock is also supported with the format 'interface:alias'`)
 		fmt.Println(`Ex: moq -pkg different . MyInterface:MyMock`)
+		fmt.Println()
+		fmt.Println(`To embed types into the mock, use the format 'interface{type1,type2...}' or 'interface:alias{type1,type2...}'`)
+		fmt.Println(`- The types must be declared in source package`)
+		fmt.Println(`- You can embed a pointer type by using the '*' prefix`)
+		fmt.Println(`Ex1: moq -pkg different src MyInterface{Type1}`)
+		fmt.Println(`Ex2: moq -pkg different src MyInterface{*Type1}`)
+		fmt.Println(`Ex3: moq -pkg different src MyInterface{Type1,Type2}`)
+		fmt.Println(`Ex4: moq -pkg different src MyInterface:MyMock{Type1}`)
 	}
 
 	flag.Parse()

pkg/moq/moq.go

@@ -6,6 +6,7 @@ import (
 	"go/token"
 	"go/types"
 	"io"
+	"regexp"
 	"strings"
 
 	"github.com/matryer/moq/internal/registry"
@@ -51,29 +52,39 @@ func New(cfg Config) (*Mocker, error) {
 }
 
 // Mock generates a mock for the specified interface name.
-func (m *Mocker) Mock(w io.Writer, namePairs ...string) error {
-	if len(namePairs) == 0 {
+func (m *Mocker) Mock(w io.Writer, targets ...string) error {
+	if len(targets) == 0 {
 		return errors.New("must specify one interface")
 	}
 
-	mocks := make([]template.MockData, len(namePairs))
-	for i, np := range namePairs {
-		name, mockName := parseInterfaceName(np)
+	differentPkg := m.registry.SrcPkgName() != m.mockPkgName()
+
+	mocks := make([]template.MockData, len(targets))
+	for i, target := range targets {
+		namePair, embeddedTypeList := parseTarget(target)
+
+		name, mockName := parseInterfaceName(namePair)
 		iface, tparams, err := m.registry.LookupInterface(name)
 		if err != nil {
 			return err
 		}
 
-		methods := make([]template.MethodData, iface.NumMethods())
+		methods := make([]template.MethodData, 0, iface.NumMethods())
 		for j := 0; j < iface.NumMethods(); j++ {
-			methods[j] = m.methodData(iface.Method(j))
+			fn := iface.Method(j)
+			if !differentPkg || fn.Exported() {
+				methods = append(methods, m.methodData(fn))
+			}
 		}
 
+		embeddedTypes := parseEmbeddedTypes(embeddedTypeList)
+
 		mocks[i] = template.MockData{
 			InterfaceName: name,
 			MockName:      mockName,
 			Methods:       methods,
 			TypeParams:    m.typeParams(tparams),
+			EmbeddedTypes: embeddedTypes,
 		}
 	}
 
@@ -88,7 +99,7 @@ func (m *Mocker) Mock(w io.Writer, namePairs ...string) error {
 	if data.MocksSomeMethod() {
 		m.registry.AddImport(types.NewPackage("sync", "sync"))
 	}
-	if m.registry.SrcPkgName() != m.mockPkgName() {
+	if differentPkg {
 		data.SrcPkgQualifier = m.registry.SrcPkgName() + "."
 		if !m.cfg.SkipEnsure {
 			imprt := m.registry.AddImport(m.registry.SrcPkg())
@@ -201,6 +212,18 @@ func (m *Mocker) format(src []byte) ([]byte, error) {
 	return gofmt(src)
 }
 
+var targetRegexp = regexp.MustCompile(`^(.+){(.+)}$`)
+
+func parseTarget(target string) (namePair, embeddedTypeList string) {
+	matched := targetRegexp.FindStringSubmatch(target)
+
+	if len(matched) != 3 {
+		return target, ""
+	}
+
+	return matched[1], matched[2]
+}
+
 func parseInterfaceName(namePair string) (ifaceName, mockName string) {
 	parts := strings.SplitN(namePair, ":", 2)
 	if len(parts) == 2 {
@@ -210,3 +233,23 @@ func parseInterfaceName(namePair string) (ifaceName, mockName string) {
 	ifaceName = parts[0]
 	return ifaceName, ifaceName + "Mock"
 }
+
+func parseEmbeddedTypes(list string) []template.EmbeddedTypeData {
+	if list == "" {
+		return nil
+	}
+
+	parts := strings.Split(list, ",")
+
+	embeddedTypes := make([]template.EmbeddedTypeData, len(parts))
+	for i, p := range parts {
+		isPointer := strings.HasPrefix(p, "*")
+		name := strings.TrimPrefix(p, "*")
+		embeddedTypes[i] = template.EmbeddedTypeData{
+			Name:      name,
+			IsPointer: isPointer,
+		}
+	}
+
+	return embeddedTypes
+}

pkg/moq/moq_test.go

@@ -418,6 +418,30 @@ func TestMockGolden(t *testing.T) {
 			interfaces: []string{"Example"},
 			goldenFile: filepath.Join("testpackages/typealias", "typealias_moq.golden.go"),
 		},
+		{
+			name:       "EmbedType",
+			cfg:        Config{SrcDir: "testpackages/embedtypes", PkgName: "mock"},
+			interfaces: []string{"Interface1{Embedded1}"},
+			goldenFile: filepath.Join("testpackages/embedtypes/mock", "embed_type.golden.go"),
+		},
+		{
+			name:       "EmbedPointerType",
+			cfg:        Config{SrcDir: "testpackages/embedtypes", PkgName: "mock"},
+			interfaces: []string{"Interface2{*Embedded2}"},
+			goldenFile: filepath.Join("testpackages/embedtypes/mock", "embed_pointer_type.golden.go"),
+		},
+		{
+			name:       "EmbedMultipleTypes",
+			cfg:        Config{SrcDir: "testpackages/embedtypes", PkgName: "mock"},
+			interfaces: []string{"Interface3{Embedded1,*Embedded2}"},
+			goldenFile: filepath.Join("testpackages/embedtypes/mock", "embed_multiple_types.golden.go"),
+		},
+		{
+			name:       "EmbedTypesWithAlias",
+			cfg:        Config{SrcDir: "testpackages/embedtypes", PkgName: "mock"},
+			interfaces: []string{"Interface3:Alias{Embedded1,*Embedded2}"},
+			goldenFile: filepath.Join("testpackages/embedtypes/mock", "embed_types_with_alias.golden.go"),
+		},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {

pkg/moq/testpackages/embedtypes/embedtypes.go

@@ -0,0 +1,36 @@
+package embedtypes
+
+// Interface1 is a test interface.
+// Its implementation must embed Embedded1.
+type Interface1 interface {
+	ExportedMethod()
+
+	unexportedMethod1()
+}
+
+// Embedded1 has unexportedMethod1.
+type Embedded1 struct{}
+
+func (e Embedded1) unexportedMethod1() {}
+
+// Interface2 is a test interface.
+// Its implementation must embed Embedded2.
+type Interface2 interface {
+	ExportedMethod()
+
+	unexportedMethod2()
+}
+
+// Embedded2 has unexportedMethod2.
+type Embedded2 struct{}
+
+func (e *Embedded2) unexportedMethod2() {}
+
+// Interface3 is a test interface.
+// Its implementation must embed Embedded1 and Embedded2.
+type Interface3 interface {
+	ExportedMethod()
+
+	unexportedMethod1()
+	unexportedMethod2()
+}