add error handle tests

rest/httpx/responses_test.go

@@ -19,13 +19,65 @@ func init() {
 }
 
 func TestError(t *testing.T) {
-	const body = "foo"
-	w := tracedResponseWriter{
-		headers: make(map[string][]string),
+	const (
+		body        = "foo"
+		wrappedBody = `"foo"`
+	)
+
+	tests := []struct {
+		name         string
+		input        string
+		errorHandler func(error) (int, interface{})
+		expectBody   string
+		expectCode   int
+	}{
+		{
+			name:       "default error handler",
+			input:      body,
+			expectBody: body,
+			expectCode: http.StatusBadRequest,
+		},
+		{
+			name:  "customized error handler return string",
+			input: body,
+			errorHandler: func(err error) (int, interface{}) {
+				return http.StatusForbidden, err.Error()
+			},
+			expectBody: wrappedBody,
+			expectCode: http.StatusForbidden,
+		},
+		{
+			name:  "customized error handler return error",
+			input: body,
+			errorHandler: func(err error) (int, interface{}) {
+				return http.StatusForbidden, err
+			},
+			expectBody: body,
+			expectCode: http.StatusForbidden,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			w := tracedResponseWriter{
+				headers: make(map[string][]string),
+			}
+			if test.errorHandler != nil {
+				lock.RLock()
+				prev := errorHandler
+				lock.RUnlock()
+				SetErrorHandler(test.errorHandler)
+				defer func() {
+					lock.Lock()
+					errorHandler = prev
+					lock.Unlock()
+				}()
+			}
+			Error(&w, errors.New(test.input))
+			assert.Equal(t, test.expectCode, w.code)
+			assert.Equal(t, test.expectBody, strings.TrimSpace(w.builder.String()))
+		})
 	}
-	Error(&w, errors.New(body))
-	assert.Equal(t, http.StatusBadRequest, w.code)
-	assert.Equal(t, body, strings.TrimSpace(w.builder.String()))
 }
 
 func TestOk(t *testing.T) {