Merge pull request #4055 from tayoogunbiyi/develop

Allow Healthcheck endpoint return JSON for Kubernetes

admin.go

@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"os"
 	"reflect"
+	"strconv"
 	"text/template"
 	"time"
 
@@ -71,7 +72,7 @@ func init() {
 // AdminIndex is the default http.Handler for admin module.
 // it matches url pattern "/".
 func adminIndex(rw http.ResponseWriter, _ *http.Request) {
-	execTpl(rw, map[interface{}]interface{}{}, indexTpl, defaultScriptsTpl)
+	writeTemplate(rw, map[interface{}]interface{}{}, indexTpl, defaultScriptsTpl)
 }
 
 // QpsIndex is the http.Handler for writing qps statistics map result info in http.ResponseWriter.
@@ -91,7 +92,7 @@ func qpsIndex(rw http.ResponseWriter, _ *http.Request) {
 		}
 	}
 
-	execTpl(rw, data, qpsTpl, defaultScriptsTpl)
+	writeTemplate(rw, data, qpsTpl, defaultScriptsTpl)
 }
 
 // ListConf is the http.Handler of displaying all beego configuration values as key/value pair.
@@ -128,7 +129,7 @@ func listConf(rw http.ResponseWriter, r *http.Request) {
 		}
 		data["Content"] = content
 		data["Title"] = "Routers"
-		execTpl(rw, data, routerAndFilterTpl, defaultScriptsTpl)
+		writeTemplate(rw, data, routerAndFilterTpl, defaultScriptsTpl)
 	case "filter":
 		var (
 			content = M{
@@ -171,7 +172,7 @@ func listConf(rw http.ResponseWriter, r *http.Request) {
 
 		data["Content"] = content
 		data["Title"] = "Filters"
-		execTpl(rw, data, routerAndFilterTpl, defaultScriptsTpl)
+		writeTemplate(rw, data, routerAndFilterTpl, defaultScriptsTpl)
 	default:
 		rw.Write([]byte("command not support"))
 	}
@@ -279,9 +280,7 @@ func profIndex(rw http.ResponseWriter, r *http.Request) {
 			http.Error(rw, err.Error(), http.StatusInternalServerError)
 			return
 		}
-
+		writeJSON(rw, dataJSON)
-		rw.Header().Set("Content-Type", "application/json")
-		rw.Write(dataJSON)
 		return
 	}
 
@@ -290,12 +289,12 @@ func profIndex(rw http.ResponseWriter, r *http.Request) {
 	if command == "gc summary" {
 		defaultTpl = gcAjaxTpl
 	}
-	execTpl(rw, data, profillingTpl, defaultTpl)
+	writeTemplate(rw, data, profillingTpl, defaultTpl)
 }
 
 // Healthcheck is a http.Handler calling health checking and showing the result.
 // it's in "/healthcheck" pattern in admin module.
-func healthcheck(rw http.ResponseWriter, _ *http.Request) {
+func healthcheck(rw http.ResponseWriter, r *http.Request) {
 	var (
 		result     []string
 		data       = make(map[interface{}]interface{})
@@ -322,10 +321,49 @@ func healthcheck(rw http.ResponseWriter, _ *http.Request) {
 		*resultList = append(*resultList, result)
 	}
 
+	queryParams := r.URL.Query()
+	jsonFlag := queryParams.Get("json")
+	shouldReturnJSON, _ := strconv.ParseBool(jsonFlag)
+
+	if shouldReturnJSON {
+		response := buildHealthCheckResponseList(resultList)
+		jsonResponse, err := json.Marshal(response)
+
+		if err != nil {
+			http.Error(rw, err.Error(), http.StatusInternalServerError)
+		} else {
+			writeJSON(rw, jsonResponse)
+		}
+		return
+	}
+
 	content["Data"] = resultList
 	data["Content"] = content
 	data["Title"] = "Health Check"
-	execTpl(rw, data, healthCheckTpl, defaultScriptsTpl)
+
+	writeTemplate(rw, data, healthCheckTpl, defaultScriptsTpl)
+}
+
+func buildHealthCheckResponseList(healthCheckResults *[][]string) []map[string]interface{} {
+	response := make([]map[string]interface{}, len(*healthCheckResults))
+
+	for i, healthCheckResult := range *healthCheckResults {
+		currentResultMap := make(map[string]interface{})
+
+		currentResultMap["name"] = healthCheckResult[0]
+		currentResultMap["message"] = healthCheckResult[1]
+		currentResultMap["status"] = healthCheckResult[2]
+
+		response[i] = currentResultMap
+	}
+
+	return response
+
+}
+
+func writeJSON(rw http.ResponseWriter, jsonData []byte) {
+	rw.Header().Set("Content-Type", "application/json")
+	rw.Write(jsonData)
 }
 
 // TaskStatus is a http.Handler with running task status (task name, status and the last execution).
@@ -371,10 +409,10 @@ func taskStatus(rw http.ResponseWriter, req *http.Request) {
 	content["Data"] = resultList
 	data["Content"] = content
 	data["Title"] = "Tasks"
-	execTpl(rw, data, tasksTpl, defaultScriptsTpl)
+	writeTemplate(rw, data, tasksTpl, defaultScriptsTpl)
 }
 
-func execTpl(rw http.ResponseWriter, data map[interface{}]interface{}, tpls ...string) {
+func writeTemplate(rw http.ResponseWriter, data map[interface{}]interface{}, tpls ...string) {
 	tmpl := template.Must(template.New("dashboard").Parse(dashboardTpl))
 	for _, tpl := range tpls {
 		tmpl = template.Must(tmpl.Parse(tpl))

admin_test.go

@@ -1,10 +1,32 @@
 package beego
 
 import (
+	"encoding/json"
+	"errors"
 	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"reflect"
+	"strings"
 	"testing"
+
+	"github.com/astaxie/beego/toolbox"
 )
 
+type SampleDatabaseCheck struct {
+}
+
+type SampleCacheCheck struct {
+}
+
+func (dc *SampleDatabaseCheck) Check() error {
+	return nil
+}
+
+func (cc *SampleCacheCheck) Check() error {
+	return errors.New("no cache detected")
+}
+
 func TestList_01(t *testing.T) {
 	m := make(M)
 	list("BConfig", BConfig, m)
@@ -75,3 +97,143 @@ func oldMap() M {
 	m["BConfig.Log.Outputs"] = BConfig.Log.Outputs
 	return m
 }
+
+func TestWriteJSON(t *testing.T) {
+	t.Log("Testing the adding of JSON to the response")
+
+	w := httptest.NewRecorder()
+	originalBody := []int{1, 2, 3}
+
+	res, _ := json.Marshal(originalBody)
+
+	writeJSON(w, res)
+
+	decodedBody := []int{}
+	err := json.NewDecoder(w.Body).Decode(&decodedBody)
+
+	if err != nil {
+		t.Fatal("Could not decode response body into slice.")
+	}
+
+	for i := range decodedBody {
+		if decodedBody[i] != originalBody[i] {
+			t.Fatalf("Expected %d but got %d in decoded body slice", originalBody[i], decodedBody[i])
+		}
+	}
+}
+
+func TestHealthCheckHandlerDefault(t *testing.T) {
+	endpointPath := "/healthcheck"
+
+	toolbox.AddHealthCheck("database", &SampleDatabaseCheck{})
+	toolbox.AddHealthCheck("cache", &SampleCacheCheck{})
+
+	req, err := http.NewRequest("GET", endpointPath, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	w := httptest.NewRecorder()
+
+	handler := http.HandlerFunc(healthcheck)
+
+	handler.ServeHTTP(w, req)
+
+	if status := w.Code; status != http.StatusOK {
+		t.Errorf("handler returned wrong status code: got %v want %v",
+			status, http.StatusOK)
+	}
+	if !strings.Contains(w.Body.String(), "database") {
+		t.Errorf("Expected 'database' in generated template.")
+	}
+
+}
+
+func TestBuildHealthCheckResponseList(t *testing.T) {
+	healthCheckResults := [][]string{
+		[]string{
+			"error",
+			"Database",
+			"Error occured whie starting the db",
+		},
+		[]string{
+			"success",
+			"Cache",
+			"Cache started successfully",
+		},
+	}
+
+	responseList := buildHealthCheckResponseList(&healthCheckResults)
+
+	if len(responseList) != len(healthCheckResults) {
+		t.Errorf("invalid response map length: got %d want %d",
+			len(responseList), len(healthCheckResults))
+	}
+
+	responseFields := []string{"name", "message", "status"}
+
+	for _, response := range responseList {
+		for _, field := range responseFields {
+			_, ok := response[field]
+			if !ok {
+				t.Errorf("expected %s to be in the response %v", field, response)
+			}
+		}
+
+	}
+
+}
+
+func TestHealthCheckHandlerReturnsJSON(t *testing.T) {
+
+	toolbox.AddHealthCheck("database", &SampleDatabaseCheck{})
+	toolbox.AddHealthCheck("cache", &SampleCacheCheck{})
+
+	req, err := http.NewRequest("GET", "/healthcheck?json=true", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	w := httptest.NewRecorder()
+
+	handler := http.HandlerFunc(healthcheck)
+
+	handler.ServeHTTP(w, req)
+	if status := w.Code; status != http.StatusOK {
+		t.Errorf("handler returned wrong status code: got %v want %v",
+			status, http.StatusOK)
+	}
+
+	decodedResponseBody := []map[string]interface{}{}
+	expectedResponseBody := []map[string]interface{}{}
+
+	expectedJSONString := []byte(`
+		[
+			{
+				"message":"database",
+				"name":"success",
+				"status":"OK"
+			},
+			{
+				"message":"cache",
+				"name":"error",
+				"status":"no cache detected"
+			}
+		]
+	`)
+
+	json.Unmarshal(expectedJSONString, &expectedResponseBody)
+
+	json.Unmarshal(w.Body.Bytes(), &decodedResponseBody)
+
+	if len(expectedResponseBody) != len(decodedResponseBody) {
+		t.Errorf("invalid response map length: got %d want %d",
+			len(decodedResponseBody), len(expectedResponseBody))
+	}
+
+	if !reflect.DeepEqual(decodedResponseBody, expectedResponseBody) {
+		t.Errorf("handler returned unexpected body: got %v want %v",
+			decodedResponseBody, expectedResponseBody)
+	}
+
+}