Skip to content

Commit a385de1

Browse files
tomMoulardtomMoulard
committed
feat: add support for generic interfaces
1 parent 2b5db55 commit a385de1

7 files changed

Lines changed: 727 additions & 42 deletions

File tree

‎Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
.PHONY: clean lint test build
1+
.PHONY: default clean lint test build
22

33
default: clean lint test build
44

‎mocktail.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,9 @@ type PackageDesc struct {
3838

3939
// InterfaceDesc represent an interface.
4040
type InterfaceDesc struct {
41-
Name string
42-
Methods []*types.Func
41+
Name string
42+
Methods []*types.Func
43+
TypeParams *types.TypeParamList // Generic type parameters
4344
}
4445

4546
func main() {
@@ -156,6 +157,11 @@ func walk(root, moduleName string) (map[string]PackageDesc, error) {
156157

157158
interfaceDesc := InterfaceDesc{Name: interfaceName}
158159

160+
// Check if this is a generic interface
161+
if namedType, ok := lookup.Type().(*types.Named); ok {
162+
interfaceDesc.TypeParams = namedType.TypeParams()
163+
}
164+
159165
interfaceType, ok := lookup.Type().Underlying().(*types.Interface)
160166
if !ok {
161167
return fmt.Errorf("type %q in %q is not an interface", lookup.Type(), fp)
@@ -255,6 +261,9 @@ func getTypeImports(t types.Type) []string {
255261
case *types.Chan:
256262
return []string{""}
257263

264+
case *types.TypeParam:
265+
return []string{""}
266+
258267
default:
259268
panic(fmt.Sprintf("OOPS %[1]T %[1]s", t))
260269
}
@@ -270,7 +279,7 @@ func generate(model map[string]PackageDesc, exported bool) error {
270279
}
271280

272281
for _, interfaceDesc := range pkgDesc.Interfaces {
273-
err = writeMockBase(buffer, interfaceDesc.Name, exported)
282+
err = writeMockBase(buffer, interfaceDesc, exported)
274283
if err != nil {
275284
return err
276285
}
@@ -285,6 +294,7 @@ func generate(model map[string]PackageDesc, exported bool) error {
285294
InterfaceName: interfaceDesc.Name,
286295
Method: method,
287296
Signature: signature,
297+
TypeParams: interfaceDesc.TypeParams,
288298
}
289299

290300
err = syrup.MockMethod(buffer)

‎syrup.go‎

Lines changed: 98 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,13 @@ package {{ .Name }}
2525

2626
templateMockBase = `
2727
// {{ .InterfaceName | ToGoCamel }}Mock mock of {{ .InterfaceName }}.
28-
type {{ .InterfaceName | ToGoCamel }}Mock struct { mock.Mock }
28+
type {{ .InterfaceName | ToGoCamel }}Mock{{ .TypeParamsDecl }} struct { mock.Mock }
2929
3030
// {{.ConstructorPrefix}}{{ .InterfaceName | ToGoPascal }}Mock creates a new {{ .InterfaceName | ToGoCamel }}Mock.
31-
func {{.ConstructorPrefix}}{{ .InterfaceName | ToGoPascal }}Mock(tb testing.TB) *{{ .InterfaceName | ToGoCamel }}Mock {
31+
func {{.ConstructorPrefix}}{{ .InterfaceName | ToGoPascal }}Mock{{ .TypeParamsDecl }}(tb testing.TB) *{{ .InterfaceName | ToGoCamel }}Mock{{ .TypeParamsUse }} {
3232
tb.Helper()
3333
34-
m := &{{ .InterfaceName | ToGoCamel }}Mock{}
34+
m := &{{ .InterfaceName | ToGoCamel }}Mock{{ .TypeParamsUse }}{}
3535
m.Mock.Test(tb)
3636
3737
tb.Cleanup(func() { m.AssertExpectations(tb) })
@@ -41,48 +41,48 @@ func {{.ConstructorPrefix}}{{ .InterfaceName | ToGoPascal }}Mock(tb testing.TB)
4141
`
4242

4343
templateCallBase = `
44-
type {{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call struct{
44+
type {{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsDecl }} struct{
4545
*mock.Call
46-
Parent *{{ .InterfaceName | ToGoCamel }}Mock
46+
Parent *{{ .InterfaceName | ToGoCamel }}Mock{{ .TypeParamsUse }}
4747
}
4848
4949
50-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) Panic(msg string) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
50+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) Panic(msg string) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
5151
_c.Call = _c.Call.Panic(msg)
5252
return _c
5353
}
5454
55-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) Once() *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
55+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) Once() *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
5656
_c.Call = _c.Call.Once()
5757
return _c
5858
}
5959
60-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) Twice() *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
60+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) Twice() *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
6161
_c.Call = _c.Call.Twice()
6262
return _c
6363
}
6464
65-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) Times(i int) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
65+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) Times(i int) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
6666
_c.Call = _c.Call.Times(i)
6767
return _c
6868
}
6969
70-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) WaitUntil(w <-chan time.Time) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
70+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) WaitUntil(w <-chan time.Time) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
7171
_c.Call = _c.Call.WaitUntil(w)
7272
return _c
7373
}
7474
75-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) After(d time.Duration) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
75+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) After(d time.Duration) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
7676
_c.Call = _c.Call.After(d)
7777
return _c
7878
}
7979
80-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) Run(fn func(args mock.Arguments)) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
80+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) Run(fn func(args mock.Arguments)) *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
8181
_c.Call = _c.Call.Run(fn)
8282
return _c
8383
}
8484
85-
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call) Maybe() *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call {
85+
func (_c *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }}) Maybe() *{{ .InterfaceName | ToGoCamel }}{{ .MethodName }}Call{{ .TypeParamsUse }} {
8686
_c.Call = _c.Call.Maybe()
8787
return _c
8888
}
@@ -96,6 +96,21 @@ type Syrup struct {
9696
InterfaceName string
9797
Method *types.Func
9898
Signature *types.Signature
99+
TypeParams *types.TypeParamList
100+
}
101+
102+
// getTypeParamsUse returns type parameters for usage in method receivers.
103+
func (s Syrup) getTypeParamsUse() string {
104+
if s.TypeParams == nil || s.TypeParams.Len() == 0 {
105+
return ""
106+
}
107+
108+
var names []string
109+
for i := range s.TypeParams.Len() {
110+
tp := s.TypeParams.At(i)
111+
names = append(names, tp.Obj().Name())
112+
}
113+
return "[" + strings.Join(names, ", ") + "]"
99114
}
100115

101116
// Call generates mock.Call wrapper.
@@ -146,7 +161,8 @@ func (s Syrup) MockMethod(writer io.Writer) error {
146161
func (s Syrup) mockedMethod(writer io.Writer) error {
147162
w := &Writer{writer: writer}
148163

149-
w.Printf("func (_m *%sMock) %s(", strcase.ToGoCamel(s.InterfaceName), s.Method.Name())
164+
typeParamsUse := s.getTypeParamsUse()
165+
w.Printf("func (_m *%sMock%s) %s(", strcase.ToGoCamel(s.InterfaceName), typeParamsUse, s.Method.Name())
150166

151167
params := s.Signature.Params()
152168

@@ -255,7 +271,8 @@ func (s Syrup) methodOn(writer io.Writer) error {
255271

256272
structBaseName := strcase.ToGoCamel(s.InterfaceName)
257273

258-
w.Printf("func (_m *%sMock) On%s(", structBaseName, s.Method.Name())
274+
typeParamsUse := s.getTypeParamsUse()
275+
w.Printf("func (_m *%sMock%s) On%s(", structBaseName, typeParamsUse, s.Method.Name())
259276

260277
params := s.Signature.Params()
261278

@@ -284,10 +301,10 @@ func (s Syrup) methodOn(writer io.Writer) error {
284301
}
285302
}
286303

287-
w.Printf(") *%s%sCall {\n", structBaseName, s.Method.Name())
304+
w.Printf(") *%s%sCall%s {\n", structBaseName, s.Method.Name(), typeParamsUse)
288305

289-
w.Printf(` return &%s%sCall{Call: _m.Mock.On("%s", %s), Parent: _m}`,
290-
structBaseName, s.Method.Name(), s.Method.Name(), strings.Join(argNames, ", "))
306+
w.Printf(` return &%s%sCall%s{Call: _m.Mock.On("%s", %s), Parent: _m}`,
307+
structBaseName, s.Method.Name(), typeParamsUse, s.Method.Name(), strings.Join(argNames, ", "))
291308

292309
w.Println()
293310
w.Println("}")
@@ -301,7 +318,8 @@ func (s Syrup) methodOnRaw(writer io.Writer) error {
301318

302319
structBaseName := strcase.ToGoCamel(s.InterfaceName)
303320

304-
w.Printf("func (_m *%sMock) On%sRaw(", structBaseName, s.Method.Name())
321+
typeParamsUse := s.getTypeParamsUse()
322+
w.Printf("func (_m *%sMock%s) On%sRaw(", structBaseName, typeParamsUse, s.Method.Name())
305323

306324
params := s.Signature.Params()
307325

@@ -330,10 +348,10 @@ func (s Syrup) methodOnRaw(writer io.Writer) error {
330348
}
331349
}
332350

333-
w.Printf(") *%s%sCall {\n", structBaseName, s.Method.Name())
351+
w.Printf(") *%s%sCall%s {\n", structBaseName, s.Method.Name(), typeParamsUse)
334352

335-
w.Printf(` return &%s%sCall{Call: _m.Mock.On("%s", %s), Parent: _m}`,
336-
structBaseName, s.Method.Name(), s.Method.Name(), strings.Join(argNames, ", "))
353+
w.Printf(` return &%s%sCall%s{Call: _m.Mock.On("%s", %s), Parent: _m}`,
354+
structBaseName, s.Method.Name(), typeParamsUse, s.Method.Name(), strings.Join(argNames, ", "))
337355

338356
w.Println()
339357
w.Println("}")
@@ -353,9 +371,26 @@ func (s Syrup) callBase(writer io.Writer) error {
353371
return err
354372
}
355373

374+
// Generate type parameter declarations and usage
375+
typeParamsDecl := ""
376+
typeParamsUse := ""
377+
if s.TypeParams != nil && s.TypeParams.Len() > 0 {
378+
var params []string
379+
var names []string
380+
for i := range s.TypeParams.Len() {
381+
tp := s.TypeParams.At(i)
382+
params = append(params, tp.Obj().Name()+" "+tp.Constraint().String())
383+
names = append(names, tp.Obj().Name())
384+
}
385+
typeParamsDecl = "[" + strings.Join(params, ", ") + "]"
386+
typeParamsUse = "[" + strings.Join(names, ", ") + "]"
387+
}
388+
356389
data := map[string]string{
357-
"InterfaceName": s.InterfaceName,
358-
"MethodName": s.Method.Name(),
390+
"InterfaceName": s.InterfaceName,
391+
"MethodName": s.Method.Name(),
392+
"TypeParamsDecl": typeParamsDecl,
393+
"TypeParamsUse": typeParamsUse,
359394
}
360395

361396
return tmpl.Execute(writer, data)
@@ -371,7 +406,8 @@ func (s Syrup) typedReturns(writer io.Writer) error {
371406

372407
structBaseName := strcase.ToGoCamel(s.InterfaceName)
373408

374-
w.Printf("func (_c *%s%sCall) TypedReturns(", structBaseName, s.Method.Name())
409+
typeParamsUse := s.getTypeParamsUse()
410+
w.Printf("func (_c *%s%sCall%s) TypedReturns(", structBaseName, s.Method.Name(), typeParamsUse)
375411

376412
var returnNames string
377413
for i := range results.Len() {
@@ -386,7 +422,7 @@ func (s Syrup) typedReturns(writer io.Writer) error {
386422
}
387423
}
388424

389-
w.Printf(") *%s%sCall {\n", structBaseName, s.Method.Name())
425+
w.Printf(") *%s%sCall%s {\n", structBaseName, s.Method.Name(), typeParamsUse)
390426
w.Printf("\t_c.Call = _c.Return(%s)\n", returnNames)
391427
w.Println("\treturn _c")
392428
w.Println("}")
@@ -402,8 +438,8 @@ func (s Syrup) typedRun(writer io.Writer) error {
402438

403439
structBaseName := strcase.ToGoCamel(s.InterfaceName)
404440

405-
w.Printf("func (_c *%[1]s%[2]sCall) TypedRun(fn %[3]s) *%[1]s%[2]sCall {\n",
406-
structBaseName, s.Method.Name(), s.createFuncSignature(params, nil))
441+
w.Printf("func (_c *%[1]s%[2]sCall%[4]s) TypedRun(fn %[3]s) *%[1]s%[2]sCall%[4]s {\n",
442+
structBaseName, s.Method.Name(), s.createFuncSignature(params, nil), s.getTypeParamsUse())
407443
w.Println("\t_c.Call = _c.Call.Run(func(args mock.Arguments) {")
408444

409445
var pos int
@@ -456,8 +492,8 @@ func (s Syrup) returnsFn(writer io.Writer) error {
456492

457493
structBaseName := strcase.ToGoCamel(s.InterfaceName)
458494

459-
w.Printf("func (_c *%[1]s%[2]sCall) ReturnsFn(fn %[3]s) *%[1]s%[2]sCall {\n",
460-
structBaseName, s.Method.Name(), s.createFuncSignature(params, results))
495+
w.Printf("func (_c *%[1]s%[2]sCall%[4]s) ReturnsFn(fn %[3]s) *%[1]s%[2]sCall%[4]s {\n",
496+
structBaseName, s.Method.Name(), s.createFuncSignature(params, results), s.getTypeParamsUse())
461497
w.Println("\t_c.Call = _c.Return(fn)")
462498
w.Println("\treturn _c")
463499
w.Println("}")
@@ -469,7 +505,8 @@ func (s Syrup) returnsFn(writer io.Writer) error {
469505
func (s Syrup) callMethodsOn(writer io.Writer, methods []*types.Func) error {
470506
w := &Writer{writer: writer}
471507

472-
callType := fmt.Sprintf("%s%sCall", strcase.ToGoCamel(s.InterfaceName), s.Method.Name())
508+
typeParamsUse := s.getTypeParamsUse()
509+
callType := fmt.Sprintf("%s%sCall%s", strcase.ToGoCamel(s.InterfaceName), s.Method.Name(), typeParamsUse)
473510

474511
for _, method := range methods {
475512
sign := method.Type().(*types.Signature)
@@ -498,7 +535,7 @@ func (s Syrup) callMethodsOn(writer io.Writer, methods []*types.Func) error {
498535
}
499536
}
500537

501-
w.Printf(") *%s%sCall {\n", strcase.ToGoCamel(s.InterfaceName), method.Name())
538+
w.Printf(") *%s%sCall%s {\n", strcase.ToGoCamel(s.InterfaceName), method.Name(), typeParamsUse)
502539

503540
w.Printf("\treturn _c.Parent.On%s(%s", method.Name(), strings.Join(argNames, ", "))
504541
if sign.Variadic() {
@@ -515,7 +552,8 @@ func (s Syrup) callMethodsOn(writer io.Writer, methods []*types.Func) error {
515552
func (s Syrup) callMethodOnRaw(writer io.Writer, methods []*types.Func) error {
516553
w := &Writer{writer: writer}
517554

518-
callType := fmt.Sprintf("%s%sCall", strcase.ToGoCamel(s.InterfaceName), s.Method.Name())
555+
typeParamsUse := s.getTypeParamsUse()
556+
callType := fmt.Sprintf("%s%sCall%s", strcase.ToGoCamel(s.InterfaceName), s.Method.Name(), typeParamsUse)
519557

520558
for _, method := range methods {
521559
sign := method.Type().(*types.Signature)
@@ -544,7 +582,7 @@ func (s Syrup) callMethodOnRaw(writer io.Writer, methods []*types.Func) error {
544582
}
545583
}
546584

547-
w.Printf(") *%s%sCall {\n", strcase.ToGoCamel(s.InterfaceName), method.Name())
585+
w.Printf(") *%s%sCall%s {\n", strcase.ToGoCamel(s.InterfaceName), method.Name(), typeParamsUse)
548586

549587
w.Printf("\treturn _c.Parent.On%sRaw(%s)\n", method.Name(), strings.Join(argNames, ", "))
550588
w.Println("}")
@@ -596,6 +634,9 @@ func (s Syrup) getTypeName(t types.Type, last bool) string {
596634
case *types.Array:
597635
return fmt.Sprintf("[%d]%s", v.Len(), s.getTypeName(v.Elem(), false))
598636

637+
case *types.TypeParam:
638+
return v.Obj().Name()
639+
599640
default:
600641
panic(fmt.Sprintf("OOPS %[1]T %[1]s", t))
601642
}
@@ -689,7 +730,7 @@ func writeImports(writer io.Writer, descPkg PackageDesc) error {
689730
return tmpl.Execute(writer, data)
690731
}
691732

692-
func writeMockBase(writer io.Writer, interfaceName string, exported bool) error {
733+
func writeMockBase(writer io.Writer, interfaceDesc InterfaceDesc, exported bool) error {
693734
base := template.New("templateMockBase").Funcs(template.FuncMap{
694735
"ToGoCamel": strcase.ToGoCamel,
695736
"ToGoPascal": strcase.ToGoPascal,
@@ -700,12 +741,33 @@ func writeMockBase(writer io.Writer, interfaceName string, exported bool) error
700741
constructorPrefix = "New"
701742
}
702743

744+
// Generate type parameter declarations and usage
745+
typeParamsDecl := ""
746+
typeParamsUse := ""
747+
if interfaceDesc.TypeParams != nil && interfaceDesc.TypeParams.Len() > 0 {
748+
var params []string
749+
var names []string
750+
for i := range interfaceDesc.TypeParams.Len() {
751+
tp := interfaceDesc.TypeParams.At(i)
752+
params = append(params, tp.Obj().Name()+" "+tp.Constraint().String())
753+
names = append(names, tp.Obj().Name())
754+
}
755+
typeParamsDecl = "[" + strings.Join(params, ", ") + "]"
756+
typeParamsUse = "[" + strings.Join(names, ", ") + "]"
757+
}
758+
703759
tmpl, err := base.Parse(templateMockBase)
704760
if err != nil {
705761
return err
706762
}
707763

708-
return tmpl.Execute(writer, map[string]interface{}{"InterfaceName": interfaceName, "ConstructorPrefix": constructorPrefix})
764+
data := map[string]interface{}{
765+
"InterfaceName": interfaceDesc.Name,
766+
"ConstructorPrefix": constructorPrefix,
767+
"TypeParamsDecl": typeParamsDecl,
768+
"TypeParamsUse": typeParamsUse,
769+
}
770+
return tmpl.Execute(writer, data)
709771
}
710772

711773
func quickGoImports(descPkg PackageDesc) []string {

‎testdata/src/a/a.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,9 @@ type Strawberry interface {
4343
type Orange interface {
4444
Juice() <-chan struct{}
4545
}
46+
47+
type Banana[T, U any] interface {
48+
Tree(T)
49+
Flower() U
50+
Pudding()
51+
}

0 commit comments

Comments
 (0)