From 6158de13efb63c6d4380583dde1522139e026774 Mon Sep 17 00:00:00 2001
From: Maneesh Babu M <maneesh@amagi.com>
Date: Fri, 23 Apr 2021 02:54:07 +0000
Subject: [PATCH] Add tests

---
 server/web/controller_test.go | 70 ++++++++++++++++++++++++++++++++++-
 1 file changed, 68 insertions(+), 2 deletions(-)

diff --git a/server/web/controller_test.go b/server/web/controller_test.go
index 8a52f097..ea7d2786 100644
--- a/server/web/controller_test.go
+++ b/server/web/controller_test.go
@@ -15,16 +15,18 @@
 package web
 
 import (
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
+	"io/ioutil"
 	"math"
 	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"strconv"
 	"testing"
 
 	"github.com/beego/beego/v2/server/web/context"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestGetInt(t *testing.T) {
@@ -244,3 +246,67 @@ func TestBindYAML(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "FOO", s.Foo)
 }
+
+type TestRespController struct {
+	Controller
+}
+
+func (t *TestRespController) TestResponse() {
+	type S struct {
+		Foo string `json:"foo" xml:"foo" yaml:"foo"`
+	}
+
+	bar := S{Foo: "bar"}
+
+	t.Resp(bar)
+}
+
+type respTestCase struct {
+	Accept                string
+	ExpectedContentLength int64
+	ExpectedResponse      string
+}
+
+func TestControllerResp(t *testing.T) {
+	// test cases
+	tcs := []respTestCase{
+		{Accept: context.ApplicationJSON, ExpectedContentLength: 13, ExpectedResponse: `{"foo":"bar"}`},
+		{Accept: context.ApplicationXML, ExpectedContentLength: 21, ExpectedResponse: `<S><foo>bar</foo></S>`},
+		{Accept: context.ApplicationYAML, ExpectedContentLength: 9, ExpectedResponse: "foo: bar\n"},
+		{Accept: "OTHER", ExpectedContentLength: 13, ExpectedResponse: `{"foo":"bar"}`},
+	}
+
+	for _, tc := range tcs {
+		testControllerRespTestCases(t, tc)
+	}
+}
+
+func testControllerRespTestCases(t *testing.T, tc respTestCase) {
+	// create fake GET request
+	r, _ := http.NewRequest("GET", "/", nil)
+	r.Header.Set("Accept", tc.Accept)
+	w := httptest.NewRecorder()
+
+	// setup the handler
+	handler := NewControllerRegister()
+	handler.Add("/", &TestRespController{}, WithRouterMethods(&TestRespController{}, "get:TestResponse"))
+	handler.ServeHTTP(w, r)
+
+	response := w.Result()
+	if response.ContentLength != tc.ExpectedContentLength {
+		t.Errorf("TestResponse() unable to validate content length for %s", tc.Accept)
+	}
+
+	if response.StatusCode != http.StatusOK {
+		t.Errorf("TestResponse() failed to validate response code for %s", tc.Accept)
+	}
+
+	bodyBytes, err := ioutil.ReadAll(response.Body)
+	if err != nil {
+		t.Errorf("TestResponse() failed to parse response body for %s", tc.Accept)
+	}
+	bodyString := string(bodyBytes)
+	if bodyString != tc.ExpectedResponse {
+		t.Errorf("TestResponse() failed to validate response body for %s", tc.Accept)
+	}
+}