From 1ceb2662ba2e07f02717fc2aa2e07b2d3cfe8003 Mon Sep 17 00:00:00 2001 From: Vishal Rana Date: Fri, 23 Sep 2016 05:31:48 -0700 Subject: [PATCH] Wrap handler and middleware functions Signed-off-by: Vishal Rana --- echo.go | 22 ++++++--- echo_test.go | 121 ++++++++++++++++++++++++++++++++++++-------------- group_test.go | 33 +++++++++++--- 3 files changed, 129 insertions(+), 47 deletions(-) diff --git a/echo.go b/echo.go index c9667de2..27237ddc 100644 --- a/echo.go +++ b/echo.go @@ -551,14 +551,22 @@ func (e *HTTPError) Error() string { return e.Message } -// WrapMiddleware wrap `echo.HandlerFunc` into `echo.MiddlewareFunc`. -func WrapMiddleware(h HandlerFunc) MiddlewareFunc { +// WrapHandler wraps `http.Handler` into `echo.HandlerFunc`. +func WrapHandler(h http.Handler) HandlerFunc { + return func(c Context) error { + h.ServeHTTP(c.Response(), c.Request()) + return nil + } +} + +// WrapMiddleware wraps `func(http.Handler) http.Handler` into `echo.MiddlewareFunc` +func WrapMiddleware(m func(http.Handler) http.Handler) MiddlewareFunc { return func(next HandlerFunc) HandlerFunc { - return func(c Context) error { - if err := h(c); err != nil { - return err - } - return next(c) + return func(c Context) (err error) { + m(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + err = next(c) + })).ServeHTTP(c.Response(), c.Request()) + return } } } diff --git a/echo_test.go b/echo_test.go index 9aade226..240b32ca 100644 --- a/echo_test.go +++ b/echo_test.go @@ -85,26 +85,34 @@ func TestEchoMiddleware(t *testing.T) { e := New() buf := new(bytes.Buffer) - e.Pre(WrapMiddleware(func(c Context) error { - assert.Empty(t, c.Path()) - buf.WriteString("-1") - return nil - })) + e.Pre(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + assert.Empty(t, c.Path()) + buf.WriteString("-1") + return next(c) + } + }) - e.Use(WrapMiddleware(func(c Context) error { - buf.WriteString("1") - return nil - })) + e.Use(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + buf.WriteString("1") + return next(c) + } + }) - e.Use(WrapMiddleware(func(c Context) error { - buf.WriteString("2") - return nil - })) + e.Use(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + buf.WriteString("2") + return next(c) + } + }) - e.Use(WrapMiddleware(func(c Context) error { - buf.WriteString("3") - return nil - })) + e.Use(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + buf.WriteString("3") + return next(c) + } + }) // Route e.GET("/", func(c Context) error { @@ -119,9 +127,11 @@ func TestEchoMiddleware(t *testing.T) { func TestEchoMiddlewareError(t *testing.T) { e := New() - e.Use(WrapMiddleware(func(c Context) error { - return errors.New("error") - })) + e.Use(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + return errors.New("error") + } + }) e.GET("/", NotFoundHandler) c, _ := request(GET, "/", e) assert.Equal(t, http.StatusInternalServerError, c) @@ -140,6 +150,43 @@ func TestEchoHandler(t *testing.T) { assert.Equal(t, "OK", b) } +func TestEchoWrapHandler(t *testing.T) { + e := New() + req, _ := http.NewRequest(GET, "", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + h := WrapHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("test")) + })) + if assert.NoError(t, h(c)) { + assert.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, "test", rec.Body.String()) + } +} + +func TestEchoWrapMiddleware(t *testing.T) { + e := New() + req, _ := http.NewRequest(GET, "", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + buf := new(bytes.Buffer) + mw := WrapMiddleware(func(h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + buf.Write([]byte("mw")) + h.ServeHTTP(w, r) + }) + }) + h := mw(func(c Context) error { + return c.String(http.StatusOK, "OK") + }) + if assert.NoError(t, h(c)) { + assert.Equal(t, "mw", buf.String()) + assert.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, "OK", rec.Body.String()) + } +} + func TestEchoConnect(t *testing.T) { e := New() testMethod(t, CONNECT, "/", e) @@ -248,10 +295,10 @@ func TestEchoRoutes(t *testing.T) { func TestEchoGroup(t *testing.T) { e := New() buf := new(bytes.Buffer) - e.Use(MiddlewareFunc(func(h HandlerFunc) HandlerFunc { + e.Use(MiddlewareFunc(func(next HandlerFunc) HandlerFunc { return func(c Context) error { buf.WriteString("0") - return h(c) + return next(c) } })) h := func(c Context) error { @@ -266,23 +313,29 @@ func TestEchoGroup(t *testing.T) { // Group g1 := e.Group("/group1") - g1.Use(WrapMiddleware(func(c Context) error { - buf.WriteString("1") - return h(c) - })) + g1.Use(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + buf.WriteString("1") + return next(c) + } + }) g1.Get("", h) // Nested groups with middleware g2 := e.Group("/group2") - g2.Use(WrapMiddleware(func(c Context) error { - buf.WriteString("2") - return nil - })) + g2.Use(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + buf.WriteString("2") + return next(c) + } + }) g3 := g2.Group("/group3") - g3.Use(WrapMiddleware(func(c Context) error { - buf.WriteString("3") - return nil - })) + g3.Use(func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + buf.WriteString("3") + return next(c) + } + }) g3.Get("", h) request(GET, "/users", e) diff --git a/group_test.go b/group_test.go index 0d131c60..e287a2b6 100644 --- a/group_test.go +++ b/group_test.go @@ -1,8 +1,9 @@ package echo import ( - "github.com/stretchr/testify/assert" "testing" + + "github.com/stretchr/testify/assert" ) // TODO: Fix me @@ -38,11 +39,31 @@ func TestGroupRouteMiddleware(t *testing.T) { e := New() g := e.Group("/group") h := func(Context) error { return nil } - m1 := WrapMiddleware(func(c Context) error { return nil }) - m2 := WrapMiddleware(func(c Context) error { return nil }) - m3 := WrapMiddleware(func(c Context) error { return nil }) - m4 := WrapMiddleware(func(c Context) error { return c.NoContent(404) }) - m5 := WrapMiddleware(func(c Context) error { return c.NoContent(405) }) + m1 := func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + return next(c) + } + } + m2 := func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + return next(c) + } + } + m3 := func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + return next(c) + } + } + m4 := func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + return c.NoContent(404) + } + } + m5 := func(next HandlerFunc) HandlerFunc { + return func(c Context) error { + return c.NoContent(405) + } + } g.Use(m1, m2, m3) g.GET("/404", h, m4) g.GET("/405", h, m5)