Merge pull request #19 from sagikazarmark/add-base-name-to-kit-endpoint

Add base name to kit endpoint generator
sagikazarmark · Dec 2, 2019 · 7697d4b · 7697d4b
2 parents 8899f46 + 210e942
commit 7697d4b
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 18 deletions.
diff --git a/internal/cmd/commands/generate/kit/endpoint.go b/internal/cmd/commands/generate/kit/endpoint.go
@@ -15,6 +15,7 @@ type endpointOptions struct {
 	serviceInterface string
 	outdir           string
 	outfile          string
+	baseName         string
 	withOc           bool
 	ocRoot           string
 }
@@ -54,6 +55,7 @@ where request and response types are any structures in the package.
 
 	flags.StringVar(&options.outdir, "outdir", "", "output directory (default: $PWD/currdir+'gen', eg. module/modulegen)")
 	flags.StringVar(&options.outfile, "outfile", "endpoint_gen.go", "output file within the output directory")
+	flags.StringVar(&options.baseName, "base-name", "", "add a base name to generated structs (default: none)")
 	flags.BoolVar(&options.withOc, "with-oc", false, "generate OpenCensus tracing middleware")
 	flags.StringVar(&options.ocRoot, "oc-root", "", "override the package name in the generated OC trace middleware")
 
@@ -120,7 +122,7 @@ func runEndpoint(options endpointOptions) error {
 
 	fmt.Printf("Generating Go kit endpoints for %s in %s\n", spec.Name, resFile)
 
-	res, err := endpoint.Generate(outpkg, spec, options.withOc, options.ocRoot)
+	res, err := endpoint.Generate(outpkg, spec, options.withOc, options.ocRoot, options.baseName)
 	if err != nil {
 		return err
 	}

diff --git a/internal/generate/kit/endpoint/generator.go b/internal/generate/kit/endpoint/generator.go
@@ -3,12 +3,13 @@ package endpoint
 import (
 	"bytes"
 	"fmt"
+	"unicode"
 
 	"github.com/dave/jennifer/jen"
 )
 
 // Generate generates a go-kit endpoint struct.
-func Generate(pkg string, spec ServiceSpec, withOc bool, ocRoot string) (string, error) {
+func Generate(pkg string, spec ServiceSpec, withOc bool, ocRoot string, baseName string) (string, error) {
 	file := jen.NewFilePath(pkg)
 
 	file.PackageComment("Code generated by mga tool. DO NOT EDIT.")
@@ -26,27 +27,42 @@ func Generate(pkg string, spec ServiceSpec, withOc bool, ocRoot string) (string,
 		)
 	}
 
-	file.Comment("Endpoints collects all of the endpoints that compose the underlying service. It's")
+	endpointStructName := "Endpoints"
+	endpointFactoryName := "MakeEndpoints"
+	endpointTraceFactoryName := "TraceEndpoints"
+	if baseName != "" {
+		// Make the first letter upper case
+		name := string(unicode.ToUpper(rune(baseName[0])))
+		if len(baseName) > 1 {
+			name += baseName[1:]
+		}
+
+		endpointStructName = fmt.Sprintf("%sEndpoints", name)
+		endpointFactoryName = fmt.Sprintf("Make%sEndpoints", name)
+		endpointTraceFactoryName = fmt.Sprintf("Trace%sEndpoints", name)
+	}
+
+	file.Commentf("%s collects all of the endpoints that compose the underlying service. It's", endpointStructName)
 	file.Comment("meant to be used as a helper struct, to collect all of the endpoints into a")
 	file.Comment("single parameter.")
-	file.Type().Id("Endpoints").Struct(endpoints...)
+	file.Type().Id(endpointStructName).Struct(endpoints...)
 
-	file.Comment("MakeEndpoints returns an Endpoints struct where each endpoint invokes")
+	file.Commentf("%s returns a(n) %s struct where each endpoint invokes", endpointFactoryName, endpointStructName)
 	file.Comment("the corresponding method on the provided service.")
-	file.Func().Id("MakeEndpoints").
+	file.Func().Id(endpointFactoryName).
 		Params(
 			jen.Id("service").Qual(spec.Package.Path, spec.Name),
 			jen.Id("middleware").Op("...").Qual("github.com/go-kit/kit/endpoint", "Middleware"),
 		).
-		Params(jen.Id("Endpoints")).
+		Params(jen.Id(endpointStructName)).
 		Block(
 			jen.Id("mw").
 				Op(":=").
 				Qual("github.com/sagikazarmark/kitx/endpoint", "Chain").
 				Call(jen.Id("middleware").Op("...")),
 			jen.Line(),
 			jen.Return(
-				jen.Id("Endpoints").Values(endpointDict),
+				jen.Id(endpointStructName).Values(endpointDict),
 			),
 		)
 
@@ -67,11 +83,15 @@ func Generate(pkg string, spec ServiceSpec, withOc bool, ocRoot string) (string,
 				Call(jen.Id("endpoints").Dot(endpoint.Name))
 		}
 
-		file.Comment("TraceEndpoints returns an Endpoints struct where each endpoint is wrapped with a tracing middleware.")
-		file.Func().Id("TraceEndpoints").
-			Params(jen.Id("endpoints").Id("Endpoints")).
-			Params(jen.Id("Endpoints")).
-			Block(jen.Return(jen.Id("Endpoints").Values(endpointDict)))
+		file.Commentf(
+			"%s returns a(n) %s struct where each endpoint is wrapped with a tracing middleware.",
+			endpointTraceFactoryName,
+			endpointStructName,
+		)
+		file.Func().Id(endpointTraceFactoryName).
+			Params(jen.Id("endpoints").Id(endpointStructName)).
+			Params(jen.Id(endpointStructName)).
+			Block(jen.Return(jen.Id(endpointStructName).Values(endpointDict)))
 	}
 
 	var buf bytes.Buffer

diff --git a/internal/generate/kit/endpoint/generator_test.go b/internal/generate/kit/endpoint/generator_test.go
@@ -23,7 +23,13 @@ func TestGenerate(t *testing.T) {
 		},
 	}
 
-	res, err := Generate("github.com/sagikazarmark/modern-go-application/internal/app/mga/todo/todogen", input, true, "")
+	res, err := Generate(
+		"github.com/sagikazarmark/modern-go-application/internal/app/mga/todo/todogen",
+		input,
+		true,
+		"",
+		"",
+	)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -46,7 +52,7 @@ type Endpoints struct {
 	OtherCall endpoint.Endpoint
 }
 
-// MakeEndpoints returns an Endpoints struct where each endpoint invokes
+// MakeEndpoints returns a(n) Endpoints struct where each endpoint invokes
 // the corresponding method on the provided service.
 func MakeEndpoints(service todo.Service, middleware ...endpoint.Middleware) Endpoints {
 	mw := kitxendpoint.Chain(middleware...)
@@ -57,7 +63,7 @@ func MakeEndpoints(service todo.Service, middleware ...endpoint.Middleware) Endp
 	}
 }
 
-// TraceEndpoints returns an Endpoints struct where each endpoint is wrapped with a tracing middleware.
+// TraceEndpoints returns a(n) Endpoints struct where each endpoint is wrapped with a tracing middleware.
 func TraceEndpoints(endpoints Endpoints) Endpoints {
 	return Endpoints{
 		Call:      kitoc.TraceEndpoint("todo.Call")(endpoints.Call),
@@ -91,6 +97,7 @@ func TestGenerate_WithOcRoot(t *testing.T) {
 		input,
 		true,
 		"path/to/pkg",
+		"",
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -114,7 +121,7 @@ type Endpoints struct {
 	OtherCall endpoint.Endpoint
 }
 
-// MakeEndpoints returns an Endpoints struct where each endpoint invokes
+// MakeEndpoints returns a(n) Endpoints struct where each endpoint invokes
 // the corresponding method on the provided service.
 func MakeEndpoints(service todo.Service, middleware ...endpoint.Middleware) Endpoints {
 	mw := kitxendpoint.Chain(middleware...)
@@ -125,7 +132,7 @@ func MakeEndpoints(service todo.Service, middleware ...endpoint.Middleware) Endp
 	}
 }
 
-// TraceEndpoints returns an Endpoints struct where each endpoint is wrapped with a tracing middleware.
+// TraceEndpoints returns a(n) Endpoints struct where each endpoint is wrapped with a tracing middleware.
 func TraceEndpoints(endpoints Endpoints) Endpoints {
 	return Endpoints{
 		Call:      kitoc.TraceEndpoint("path/to/pkg.Call")(endpoints.Call),
@@ -136,3 +143,72 @@ func TraceEndpoints(endpoints Endpoints) Endpoints {
 
 	assert.Equal(t, expected, res, "the generated code does not match the expected one")
 }
+
+func TestGenerate_WithBase_name(t *testing.T) {
+	input := ServiceSpec{
+		Name: "Service",
+		Package: PackageSpec{
+			Name: "todo",
+			Path: "github.com/sagikazarmark/modern-go-application/internal/app/mga/todo",
+		},
+		Endpoints: []EndpointSpec{
+			{
+				Name: "Call",
+			},
+			{
+				Name: "OtherCall",
+			},
+		},
+	}
+
+	res, err := Generate(
+		"github.com/sagikazarmark/modern-go-application/internal/app/mga/todo/todogen",
+		input,
+		true,
+		"",
+		"Todo",
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected := `// Code generated by mga tool. DO NOT EDIT.
+package todogen
+
+import (
+	"github.com/go-kit/kit/endpoint"
+	kitoc "github.com/go-kit/kit/tracing/opencensus"
+	kitxendpoint "github.com/sagikazarmark/kitx/endpoint"
+	"github.com/sagikazarmark/modern-go-application/internal/app/mga/todo"
+)
+
+// TodoEndpoints collects all of the endpoints that compose the underlying service. It's
+// meant to be used as a helper struct, to collect all of the endpoints into a
+// single parameter.
+type TodoEndpoints struct {
+	Call      endpoint.Endpoint
+	OtherCall endpoint.Endpoint
+}
+
+// MakeTodoEndpoints returns a(n) TodoEndpoints struct where each endpoint invokes
+// the corresponding method on the provided service.
+func MakeTodoEndpoints(service todo.Service, middleware ...endpoint.Middleware) TodoEndpoints {
+	mw := kitxendpoint.Chain(middleware...)
+
+	return TodoEndpoints{
+		Call:      mw(MakeCallEndpoint(service)),
+		OtherCall: mw(MakeOtherCallEndpoint(service)),
+	}
+}
+
+// TraceTodoEndpoints returns a(n) TodoEndpoints struct where each endpoint is wrapped with a tracing middleware.
+func TraceTodoEndpoints(endpoints TodoEndpoints) TodoEndpoints {
+	return TodoEndpoints{
+		Call:      kitoc.TraceEndpoint("todo.Call")(endpoints.Call),
+		OtherCall: kitoc.TraceEndpoint("todo.OtherCall")(endpoints.OtherCall),
+	}
+}
+`
+
+	assert.Equal(t, expected, res, "the generated code does not match the expected one")
+}