From a6dd404d60492c0f51b0de026c597a779db66be0 Mon Sep 17 00:00:00 2001
From: Deardrops <passagesfall@gmail.com>
Date: Sat, 29 Jun 2019 15:35:28 +0800
Subject: [PATCH] code: improve wrapper & adatper test coverage

---
 wrapper.go      |   4 +-
 wrapper_test.go | 112 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 113 insertions(+), 3 deletions(-)

diff --git a/wrapper.go b/wrapper.go
index a13699a..1c6e685 100644
--- a/wrapper.go
+++ b/wrapper.go
@@ -57,11 +57,11 @@ func failure(w http.ResponseWriter, err error) {
 		statusCode = v.StatusCode()
 	}
 	w.WriteHeader(statusCode)
-	json.NewEncoder(w).Encode(errorEncoder(err))
+	json.NewEncoder(w).Encode(errorEncoder(nil, err))
 }
 
 func success(w http.ResponseWriter, data interface{}) {
-	json.NewEncoder(w).Encode(responseEncoder(data))
+	json.NewEncoder(w).Encode(responseEncoder(nil, data))
 }
 
 func (fn *fn) ServeHTTP(w http.ResponseWriter, r *http.Request) {
diff --git a/wrapper_test.go b/wrapper_test.go
index 714d341..e70348a 100644
--- a/wrapper_test.go
+++ b/wrapper_test.go
@@ -16,12 +16,15 @@ package fn
 import (
 	"bytes"
 	"context"
+	"encoding/json"
+	"errors"
 	"io"
 	"io/ioutil"
 	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"reflect"
 	"testing"
 
 	. "github.com/pingcap/check"
@@ -29,7 +32,7 @@ import (
 
 type fnSuite struct{}
 
-func Testfn(t *testing.T) {
+func TestFn(t *testing.T) {
 	TestingT(t)
 }
 
@@ -43,6 +46,10 @@ type testResponse struct {
 	Code    int    `json:"code"`
 	Message string `json:"message"`
 }
+type testErrorResponse struct {
+	Code  int    `json:"code"`
+	Error string `json:"message"`
+}
 
 var successResponse = &testResponse{Message: "success"}
 
@@ -129,6 +136,109 @@ func (s *fnSuite) TestPlugin(c *C) {
 	handler.ServeHTTP(recorder, request)
 }
 
+func (s *fnSuite) TestGroupPlugin(c *C) {
+	group := NewGroup()
+	group.Plugin(func(ctx context.Context, request *http.Request) (context.Context, error) {
+		return context.WithValue(ctx, "key", "value"), nil
+	})
+
+	logic := func(ctx context.Context) (*testResponse, error) {
+		c.Assert(ctx.Value("key").(string) == "value", IsTrue)
+		return &testResponse{}, nil
+	}
+	handler := group.Wrap(logic)
+
+	recorder := httptest.NewRecorder()
+	request, err := http.NewRequest(http.MethodGet, "", nil)
+	c.Assert(err, IsNil)
+	handler.ServeHTTP(recorder, request)
+}
+
+func (s *fnSuite) TestSetResponseEncoder(c *C) {
+	handler := Wrap(func(ctx context.Context, request *http.Request) (context.Context, error) {
+		return nil, nil
+	})
+
+	testResp := &testResponse{
+		Code:    1,
+		Message: "msg",
+	}
+	SetResponseEncoder(func(ctx context.Context, payload interface{}) interface{} {
+		return testResp
+	})
+
+	recorder := httptest.NewRecorder()
+	request, err := http.NewRequest(http.MethodGet, "", nil)
+	c.Assert(err, IsNil)
+	handler.ServeHTTP(recorder, request)
+	respMsg := &testResponse{}
+	json.Unmarshal(recorder.Body.Bytes(), &respMsg)
+	c.Assert(reflect.DeepEqual(respMsg, testResp), IsTrue)
+}
+
+func (s *fnSuite) TestSetErrorEncoder(c *C) {
+	handler := Wrap(func(ctx context.Context, request *http.Request) (context.Context, error) {
+		return nil, errors.New("")
+	})
+
+	testErrorResp := &testErrorResponse{
+		Code:  -1,
+		Error: "something went wrong",
+	}
+	SetErrorEncoder(func(ctx context.Context, err error) interface{} {
+		return testErrorResp
+	})
+
+	recorder := httptest.NewRecorder()
+	request, err := http.NewRequest(http.MethodGet, "", nil)
+	c.Assert(err, IsNil)
+	handler.ServeHTTP(recorder, request)
+
+	respMsg := &testErrorResponse{}
+	json.Unmarshal(recorder.Body.Bytes(), &respMsg)
+	c.Assert(reflect.DeepEqual(respMsg, testErrorResp), IsTrue)
+}
+
+func (s *fnSuite) TestGenericAdapter_Invoke(c *C) {
+	type CustomForm testRequest
+	handler := Wrap(func(ctx context.Context, form *CustomForm) (context.Context, error) {
+		return nil, nil
+	})
+
+	recorder := httptest.NewRecorder()
+	request, err := http.NewRequest(http.MethodGet, "", nil)
+	payload := []byte(`{"for":"hello", "bar":10000}`)
+	request.Body = ioutil.NopCloser(bytes.NewBuffer(payload))
+	c.Assert(err, IsNil)
+	handler.ServeHTTP(recorder, request)
+}
+
+func (s *fnSuite) TestSimpleUnaryAdapter_Invoke(c *C) {
+	handler := Wrap(withReq)
+
+	recorder := httptest.NewRecorder()
+	request, err := http.NewRequest(http.MethodGet, "", nil)
+	if err != nil {
+		c.Fatal(err)
+	}
+	payload := []byte(`{"for":"hello", "bar":10000}`)
+	request.Body = ioutil.NopCloser(bytes.NewBuffer(payload))
+	c.Assert(err, IsNil)
+	handler.ServeHTTP(recorder, request)
+}
+
+func (s *fnSuite) TestErrorWithStatusCode(c *C) {
+	handler := Wrap(func(ctx context.Context, request *http.Request) (context.Context, error) {
+		return nil, ErrorWithStatusCode(errors.New("not found"), http.StatusNotFound)
+	})
+
+	recorder := httptest.NewRecorder()
+	request, err := http.NewRequest(http.MethodGet, "", nil)
+	c.Assert(err, IsNil)
+	handler.ServeHTTP(recorder, request)
+	c.Assert(recorder.Code == http.StatusNotFound, IsTrue)
+}
+
 func BenchmarkSimplePlainAdapter_Invoke(b *testing.B) {
 	handler := Wrap(withNone)
 	request, err := http.NewRequest(http.MethodGet, "", nil)