From eb5ee489d9bbbd5085538525db12f0e13148748e Mon Sep 17 00:00:00 2001
From: Alexander Nevskiy <kepkin@gmail.com>
Date: Wed, 13 May 2020 14:44:38 +0300
Subject: [PATCH] Add custom router to register middleware.

---
 internal/generator/router.go | 14 ++++++++++++++
 test/api/api_gorest.go       | 18 ++++++++++++++++++
 test/generated_api_test.go   | 26 ++++++++++++++++++++++++++
 3 files changed, 58 insertions(+)

diff --git a/internal/generator/router.go b/internal/generator/router.go
index edf71be..a18072c 100644
--- a/internal/generator/router.go
+++ b/internal/generator/router.go
@@ -24,6 +24,20 @@ func RegisterRoutes(r *gin.Engine, api {{ .InterfaceName }}) {
             {{- end }}
         {{- end }}
 	{{ end -}}
+}
+
+type HandlerRegister func(operationID, httpMethod, relativePath string, handler gin.HandlerFunc)
+
+// CustomRouter
+func RegisterRoutesCustom(handlerRegister HandlerRegister, api {{ .InterfaceName }}) {
+	e := &{{ .InterfaceName }}Server{api}
+	{{ range $url, $m := .Paths -}}
+        {{- range $methodName, $method := (dict "GET" $m.Get "POST" $m.Post "PATCH" $m.Patch "DELETE" $m.Delete "PUT" $m.Put "OPTIONS" $m.Options) }}
+			{{- with $method }}
+				handlerRegister("{{ .OperationID }}", "{{ $methodName }}", "{{ ConvertUrl $url }}", e._{{ $.InterfaceName }}_{{ .OperationID }}_Handler)
+            {{- end }}
+        {{- end }}
+	{{ end -}}
 }`))
 
 func (Generator) makeRouter(wr io.Writer, sp openapi3.Spec) error {
diff --git a/test/api/api_gorest.go b/test/api/api_gorest.go
index 637de94..834ab37 100644
--- a/test/api/api_gorest.go
+++ b/test/api/api_gorest.go
@@ -643,6 +643,24 @@ func RegisterRoutes(r *gin.Engine, api PaymentGatewayAPI) {
 	r.Handle("GET", "/v1/user/:userId", e._PaymentGatewayAPI_GetUser_Handler)
 }
 
+type HandlerRegister func(operationID, httpMethod, relativePath string, handler gin.HandlerFunc)
+
+// CustomRouter
+func RegisterRoutesCustom(handlerRegister HandlerRegister, api PaymentGatewayAPI) {
+	e := &PaymentGatewayAPIServer{api}
+
+	handlerRegister("Example", "GET", "/v1/example/:year/:user", e._PaymentGatewayAPI_Example_Handler)
+
+	handlerRegister("GetFile", "GET", "/v1/files/:filename", e._PaymentGatewayAPI_GetFile_Handler)
+
+	handlerRegister("GetPayment", "GET", "/v1/payment", e._PaymentGatewayAPI_GetPayment_Handler)
+	handlerRegister("ProvidePayment", "POST", "/v1/payment", e._PaymentGatewayAPI_ProvidePayment_Handler)
+
+	handlerRegister("CreateUser", "POST", "/v1/user", e._PaymentGatewayAPI_CreateUser_Handler)
+
+	handlerRegister("GetUser", "GET", "/v1/user/:userId", e._PaymentGatewayAPI_GetUser_Handler)
+}
+
 type ID string
 
 func (t ID) Validate() (errors []FieldError) {
diff --git a/test/generated_api_test.go b/test/generated_api_test.go
index d5fb053..8dd394d 100644
--- a/test/generated_api_test.go
+++ b/test/generated_api_test.go
@@ -116,3 +116,29 @@ func TestCreateUser(t *testing.T) {
 
 	assert.Equal(t, response.Body.Bytes(), photoBytes)
 }
+
+
+func TestCustomRouteWithMiddleWare(t *testing.T) {
+	r := gin.New()
+	api.RegisterRoutesCustom(func(operationID, httpMethod, relativePath string, handler gin.HandlerFunc) {
+		handlers := make([]gin.HandlerFunc, 0)
+
+		if operationID == "GetUser" {
+			handlers = append(handlers, func(c *gin.Context){
+				c.AbortWithStatus(401)
+			})
+		}
+
+		handlers = append(handlers, handler)
+		r.Handle(httpMethod, relativePath, handlers...)
+	}, api.NewPaymentGatewayAPI())
+
+	// Go to route with middle ware
+	request := httptest.NewRequest(http.MethodGet, "/v1/user/123", nil)
+	response := httptest.NewRecorder()
+
+	r.ServeHTTP(response, request)
+	if !assert.Equal(t, http.StatusUnauthorized, response.Code, response.Body.String()) {
+		return
+	}
+}