Generate handlers in tchannel client test helper

uber · Apr 20, 2017 · 8cb587e · 8cb587e
1 parent 39bfe64
commit 8cb587e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 49 deletions.
diff --git a/codegen/template.go b/codegen/template.go
@@ -98,6 +98,7 @@ type ClientStub struct {
 
 var camelingRegex = regexp.MustCompile("[0-9A-Za-z]+")
 var funcMap = tmpl.FuncMap{
+	"lower":        strings.ToLower,
 	"title":        strings.Title,
 	"Title":        strings.Title,
 	"fullTypeName": fullTypeName,

diff --git a/codegen/templates/tchannel_client_test_server.tmpl b/codegen/templates/tchannel_client_test_server.tmpl
@@ -18,59 +18,43 @@ import (
 	{{end}}
 )
 
-{{$clientID := .ClientID -}}
-{{$interface := title .ClientID | printf "TChan%s" -}}
-{{$clientName := title .ClientID | printf "%sClient" -}}
 {{range $svc := .Services}}
-{{$server := $svc.Name | printf "%sServer"}}
-// {{$server}} is the TChannel backend for {{$svc.Name}} service.
-type {{$server}} struct {
-    handler {{$interface}}
-}
+{{range .Methods}}
+{{$privateName := lower .Name -}}
+{{$genCodePkg := .GenCodePkgName -}}
+{{$func := printf "%s%sFunc" $svc.Name .Name -}}
+{{$handler := printf "%s%sHandler" $svc.Name .Name -}}
+// {{$func}} is the handler function for "{{.Name}}" method of thrift service "{{$svc.Name}}".
+type {{$func}} func (
+    ctx context.Context,
+    reqHeaders map[string]string,
+    {{if ne .RequestType "" -}}
+        args *{{.RequestType}},
+    {{end -}}
+) ({{- if ne .ResponseType "" -}}*{{.ResponseType}}, {{- end -}}map[string]string, error)
 
-// New{{$server}} wraps a handler for {{title $clientID}} so it can be registered with a thrift server.
-func New{{$server}}(handler {{$interface}}) zanzibar.TChanServer {
-	return &{{$server}}{
-	    handler,
-	}
+// New{{$handler}} wraps a handler function so it can be registered with a thrift server.
+func New{{$handler}}(f {{$func}}) zanzibar.TChanHandler {
+    return &{{$handler}}{f}
 }
 
-// Service returns the service name.
-func (s *{{$server}}) Service() string {
-	return "{{$svc.Name}}"
+// {{$handler}} handles the "{{.Name}}" method call of thrift service "{{$svc.Name}}".
+type {{$handler}} struct {
+    {{$privateName}} {{$func}}
 }
 
-// Methods returns the method names handled by this server.
-func (s *{{$server}}) Methods() []string {
-	return []string{
-	{{range .Methods -}}
-	    "{{.Name}}",
-	{{end -}}
-	}
+// Service returns the service name.
+func (h *{{$handler}}) Service() string {
+    return "{{$svc.Name}}"
 }
 
-// Handle dispatches a method call to corresponding handler.
-func (s *{{$server}}) Handle(
-    ctx context.Context,
-    methodName string,
-    reqHeaders map[string]string,
-    wireValue *wire.Value,
-) (bool, zanzibar.RWTStruct, map[string]string, error) {
-	switch methodName {
-	{{range .Methods -}}
-	    case "{{.Name}}":
-	        return s.handle{{.Name}}(ctx, reqHeaders, wireValue)
-	{{end -}}
-	default:
-	    return false, nil, nil, fmt.Errorf(
-	        "method %v not found in service %v", methodName, s.Service(),
-	    )
-	}
+// Method returns the method name.
+func (h *{{$handler}}) Method() string {
+    return "{{.Name}}"
 }
 
-{{range .Methods}}
-{{$genCodePkg := .GenCodePkgName -}}
-func (s *{{$server}}) handle{{.Name}}(
+// Handle parses request from wire value and calls corresponding handler function.
+func (h *{{$handler}}) Handle(
     ctx context.Context,
     reqHeaders map[string]string,
     wireValue *wire.Value,
@@ -83,28 +67,29 @@ func (s *{{$server}}) handle{{.Name}}(
 	}
 
     {{- if and (eq .RequestType "") (eq .ResponseType "")}}
-	    respHeaders, err := s.handler.{{.Name}}(ctx, reqHeaders)
+	    respHeaders, err := h.{{$privateName}}(ctx, reqHeaders)
     {{- else if eq .RequestType ""}}
-	    r, respHeaders, err := s.handler.{{.Name}}(ctx, reqHeaders)
+	    r, respHeaders, err := h.{{$privateName}}(ctx, reqHeaders)
     {{- else if eq .ResponseType ""}}
-	    respHeaders, err := s.handler.{{.Name}}(ctx, reqHeaders, &req)
+	    respHeaders, err := h.{{$privateName}}(ctx, reqHeaders, &req)
 	{{- else}}
-	    r, respHeaders, err := s.handler.{{.Name}}(ctx, reqHeaders, &req)
+	    r, respHeaders, err := h.{{$privateName}}(ctx, reqHeaders, &req)
 	{{- end}}
 
-	{{if eq (len .Exceptions) 0 -}}
+    {{if eq (len .Exceptions) 0 -}}
 	    if err != nil {
             return false, nil, nil, err
 	    }
 	    res.Success = r
     {{else -}}
 	    if err != nil {
 	        switch v := err.(type) {
+	        {{$method := .Name -}}
 	        {{range .Exceptions -}}
 	            case *{{$genCodePkg}}.{{.Name}}:
 	                if v == nil {
 	                    return false, nil, nil, errors.New(
-	                        "Handler for {{.Name}} returned non-nil error type *{{.Name}} but nil value",
+	                        "Handler for {{$method}} returned non-nil error type *{{.Name}} but nil value",
 	                    )
 	                }
 	                res.{{.Name}} = v