Set content-type for mixed-in Prometheus results

Signed-off-by: Alex Ellis <alexellis2@gmail.com>

gateway/metrics/add_metrics.go

@@ -18,10 +18,8 @@ func makeClient() http.Client {
 }
 
 // AddMetricsHandler wraps a http.HandlerFunc with Prometheus metrics
-func AddMetricsHandler(handler http.HandlerFunc, host string, port int) http.HandlerFunc {
-	client := makeClient()
+func AddMetricsHandler(handler http.HandlerFunc, prometheusQuery PrometheusQueryFetcher) http.HandlerFunc {
 
-	prometheusQuery := NewPrometheusQuery(host, port, &client)
 	return func(w http.ResponseWriter, r *http.Request) {
 		// log.Printf("Calling upstream for function info\n")
 
@@ -30,11 +28,14 @@ func AddMetricsHandler(handler http.HandlerFunc, host string, port int) http.Han
 		upstreamCall := recorder.Result()
 
 		if upstreamCall.Body == nil {
+			log.Println("Upstream call had empty body.")
 			return
 		}
+
 		defer upstreamCall.Body.Close()
 
 		if recorder.Code != http.StatusOK {
+			w.Header().Set("Content-Type", "text/plain")
 			w.WriteHeader(http.StatusInternalServerError)
 			w.Write([]byte(fmt.Sprintf("Error pulling metrics from provider/backend. Status code: %d", recorder.Code)))
 			return
@@ -44,9 +45,11 @@ func AddMetricsHandler(handler http.HandlerFunc, host string, port int) http.Han
 		var functions []requests.Function
 
 		err := json.Unmarshal(upstreamBody, &functions)
+
 		if err != nil {
 			log.Println(err)
 
+			w.Header().Set("Content-Type", "text/plain")
 			w.WriteHeader(http.StatusInternalServerError)
 			w.Write([]byte("Error parsing metrics from upstream provider/backend."))
 			return
@@ -58,7 +61,7 @@ func AddMetricsHandler(handler http.HandlerFunc, host string, port int) http.Han
 		results, fetchErr := prometheusQuery.Fetch(expr)
 		if fetchErr != nil {
 			log.Printf("Error querying Prometheus API: %s\n", fetchErr.Error())
-
+			w.Header().Set("Content-Type", "application/json")
 			w.WriteHeader(http.StatusOK)
 			w.Write(upstreamBody)
 			return
@@ -73,6 +76,7 @@ func AddMetricsHandler(handler http.HandlerFunc, host string, port int) http.Han
 		}
 
 		// log.Printf("Writing bytesOut: %s\n", bytesOut)
+		w.Header().Set("Content-Type", "application/json")
 		w.WriteHeader(http.StatusOK)
 		w.Write(bytesOut)
 	}

gateway/metrics/prometheus_query.go

@@ -14,6 +14,10 @@ type PrometheusQuery struct {
 	Client *http.Client
 }
 
+type PrometheusQueryFetcher interface {
+	Fetch(query string) (*VectorQueryResponse, error)
+}
+
 // NewPrometheusQuery create a NewPrometheusQuery
 func NewPrometheusQuery(host string, port int, client *http.Client) PrometheusQuery {
 	return PrometheusQuery{
@@ -24,7 +28,7 @@ func NewPrometheusQuery(host string, port int, client *http.Client) PrometheusQu
 }
 
 // Fetch queries aggregated stats
-func (q *PrometheusQuery) Fetch(query string) (*VectorQueryResponse, error) {
+func (q PrometheusQuery) Fetch(query string) (*VectorQueryResponse, error) {
 
 	req, reqErr := http.NewRequest("GET", fmt.Sprintf("http://%s:%d/api/v1/query/?query=%s", q.Host, q.Port, query), nil)
 	if reqErr != nil {

gateway/server.go

@@ -120,7 +120,8 @@ func main() {
 		faasHandlers.AsyncReport = internalHandlers.MakeAsyncReport(metricsOptions)
 	}
 
-	listFunctions := metrics.AddMetricsHandler(faasHandlers.ListFunctions, config.PrometheusHost, config.PrometheusPort)
+	prometheusQuery := metrics.NewPrometheusQuery(config.PrometheusHost, config.PrometheusPort, &http.Client{})
+	listFunctions := metrics.AddMetricsHandler(faasHandlers.ListFunctions, prometheusQuery)
 
 	r := mux.NewRouter()
 

gateway/tests/add_metrics_test.go

@@ -0,0 +1,94 @@
+package tests
+
+import (
+	"encoding/json"
+	"log"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/openfaas/faas/gateway/metrics"
+	"github.com/openfaas/faas/gateway/requests"
+)
+
+type FakePrometheusQueryFetcher struct {
+}
+
+func (q FakePrometheusQueryFetcher) Fetch(query string) (*metrics.VectorQueryResponse, error) {
+
+	return &metrics.VectorQueryResponse{}, nil
+}
+
+func makeFakePrometheusQueryFetcher() FakePrometheusQueryFetcher {
+	return FakePrometheusQueryFetcher{}
+}
+
+func Test_PrometheusMetrics_MixedInto_Services(t *testing.T) {
+	functionsHandler := makeFunctionsHandler()
+	fakeQuery := makeFakePrometheusQueryFetcher()
+
+	handler := metrics.AddMetricsHandler(functionsHandler, fakeQuery)
+
+	rr := httptest.NewRecorder()
+	request, _ := http.NewRequest(http.MethodGet, "/system/functions", nil)
+	handler.ServeHTTP(rr, request)
+
+	if status := rr.Code; status != http.StatusOK {
+		t.Errorf("handler returned wrong status code: got %v want %v", status, http.StatusOK)
+	}
+
+	if rr.Header().Get("Content-Type") != "application/json" {
+		t.Errorf("Want application/json content-type, got: %s", rr.Header().Get("Content-Type"))
+	}
+	if len(rr.Body.String()) == 0 {
+		t.Errorf("Want content-length > 0, got: %d", len(rr.Body.String()))
+	}
+
+}
+
+func Test_FunctionsHandler_ReturnsJSONAndOneFunction(t *testing.T) {
+	functionsHandler := makeFunctionsHandler()
+
+	rr := httptest.NewRecorder()
+	request, err := http.NewRequest(http.MethodGet, "/system/functions", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	functionsHandler.ServeHTTP(rr, request)
+
+	if status := rr.Code; status != http.StatusOK {
+		t.Errorf("handler returned wrong status code: got %v want %v", status, http.StatusOK)
+	}
+
+	if rr.Header().Get("Content-Type") != "application/json" {
+		t.Errorf("Want application/json content-type, got: %s", rr.Header().Get("Content-Type"))
+	}
+
+	if len(rr.Body.String()) == 0 {
+		t.Errorf("Want content-length > 0, got: %d", len(rr.Body.String()))
+	}
+
+}
+
+func makeFunctionsHandler() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		functions := []requests.Function{
+			requests.Function{
+				Name:     "echo",
+				Replicas: 0,
+			},
+		}
+		bytesOut, marshalErr := json.Marshal(&functions)
+		if marshalErr != nil {
+			log.Fatal(marshalErr.Error())
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		_, err := w.Write(bytesOut)
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+}