Add unit tests for all the /json cases

Signed-off-by: David Weitzman <[email protected]>

src/server/server_impl.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"bytes"
 	"expvar"
 	"fmt"
 	"io"
@@ -53,15 +54,15 @@ func (server *server) AddDebugHttpEndpoint(path string, help string, handler htt
 	server.debugListener.endpoints[path] = help
 }
 
-// add an http/1 handler at the /json endpoint which allows this ratelimit service to work with
+// create an http/1 handler at the /json endpoint which allows this ratelimit service to work with
 // clients that cannot use the gRPC interface (e.g. lua)
 // example usage from cURL with domain "dummy" and descriptor "perday":
 // echo '{"domain": "dummy", "descriptors": [{"entries": [{"key": "perday"}]}]}' | curl -vvvXPOST --data @/dev/stdin localhost:8080/json
-func (server *server) AddJsonHandler(svc pb.RateLimitServiceServer) {
+func NewJsonHandler(svc pb.RateLimitServiceServer) func(http.ResponseWriter, *http.Request) {
 	// Default options include enums as strings and no identation.
 	m := &jsonpb.Marshaler{}
 
-	handler := func(writer http.ResponseWriter, request *http.Request) {
+	return func(writer http.ResponseWriter, request *http.Request) {
 		var req pb.RateLimitRequest
 
 		if err := jsonpb.Unmarshal(request.Body, &req); err != nil {
@@ -79,21 +80,26 @@ func (server *server) AddJsonHandler(svc pb.RateLimitServiceServer) {
 
 		logger.Debugf("resp:%s", resp)
 
-		writer.Header().Set("Content-Type", "application/json")
-		if resp.OverallCode == pb.RateLimitResponse_OVER_LIMIT {
-			writer.WriteHeader(http.StatusTooManyRequests)
-		} else if resp.OverallCode == pb.RateLimitResponse_UNKNOWN {
-			writer.WriteHeader(http.StatusInternalServerError)
-		}
-
-		err = m.Marshal(writer, resp)
+		buf := bytes.NewBuffer(nil)
+		err = m.Marshal(buf, resp)
 		if err != nil {
-			writer.Header().Set("Content-Type", "text/plain")
-			fmt.Fprintf(writer, "Internal error marshaling proto3 to json: %v", err)
+			logger.Errorf("error marshaling proto3 to json: %s", err.Error())
+			http.Error(writer, "error marshaling proto3 to json: "+err.Error(), http.StatusInternalServerError)
+			return
 		}
 
+		writer.Header().Set("Content-Type", "application/json")
+		if resp == nil || resp.OverallCode == pb.RateLimitResponse_UNKNOWN {
+			writer.WriteHeader(http.StatusInternalServerError)
+		} else if resp.OverallCode == pb.RateLimitResponse_OVER_LIMIT {
+			writer.WriteHeader(http.StatusTooManyRequests)
+		}
+		writer.Write(buf.Bytes())
 	}
-	server.router.HandleFunc("/json", handler)
+}
+
+func (server *server) AddJsonHandler(svc pb.RateLimitServiceServer) {
+	server.router.HandleFunc("/json", NewJsonHandler(svc))
 }
 
 func (server *server) GrpcServer() *grpc.Server {

test/mocks/mocks.go

@@ -5,3 +5,4 @@ package mocks
 //go:generate go run github.com/golang/mock/mockgen -destination ./config/config.go github.com/envoyproxy/ratelimit/src/config RateLimitConfig,RateLimitConfigLoader
 //go:generate go run github.com/golang/mock/mockgen -destination ./redis/redis.go github.com/envoyproxy/ratelimit/src/redis Client
 //go:generate go run github.com/golang/mock/mockgen -destination ./limiter/limiter.go github.com/envoyproxy/ratelimit/src/limiter RateLimitCache,TimeSource,JitterRandSource
+//go:generate go run github.com/golang/mock/mockgen -destination ./rls/rls.go github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v2 RateLimitServiceServer

test/server/server_impl_test.go

@@ -0,0 +1,86 @@
+package server_test
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v2"
+
+	"github.com/envoyproxy/ratelimit/src/server"
+	mock_v2 "github.com/envoyproxy/ratelimit/test/mocks/rls"
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+)
+
+func assertHttpResponse(t *testing.T,
+	handler http.HandlerFunc,
+	requestBody string,
+	expectedStatusCode int,
+	expectedContentType string,
+	expectedResponseBody string) {
+
+	t.Helper()
+	assert := assert.New(t)
+
+	req := httptest.NewRequest("METHOD_NOT_CHECKED", "/path_not_checked", strings.NewReader(requestBody))
+	w := httptest.NewRecorder()
+	handler(w, req)
+
+	resp := w.Result()
+	actualBody, _ := ioutil.ReadAll(resp.Body)
+	assert.Equal(expectedContentType, resp.Header.Get("Content-Type"))
+	assert.Equal(expectedStatusCode, resp.StatusCode)
+	assert.Equal(expectedResponseBody, string(actualBody))
+}
+
+func TestJsonHandler(t *testing.T) {
+	controller := gomock.NewController(t)
+	defer controller.Finish()
+
+	rls := mock_v2.NewMockRateLimitServiceServer(controller)
+	handler := server.NewJsonHandler(rls)
+
+	// Missing request body
+	assertHttpResponse(t, handler, "", 400, "text/plain; charset=utf-8", "EOF\n")
+
+	// Request body is not valid json
+	assertHttpResponse(t, handler, "}", 400, "text/plain; charset=utf-8", "invalid character '}' looking for beginning of value\n")
+
+	// Unknown response code
+	rls.EXPECT().ShouldRateLimit(nil, &pb.RateLimitRequest{
+		Domain: "foo",
+	}).Return(&pb.RateLimitResponse{}, nil)
+	assertHttpResponse(t, handler, `{"domain": "foo"}`, 500, "application/json", "{}")
+
+	// ratelimit service error
+	rls.EXPECT().ShouldRateLimit(nil, &pb.RateLimitRequest{
+		Domain: "foo",
+	}).Return(nil, fmt.Errorf("some error"))
+	assertHttpResponse(t, handler, `{"domain": "foo"}`, 400, "text/plain; charset=utf-8", "some error\n")
+
+	// json unmarshaling error
+	rls.EXPECT().ShouldRateLimit(nil, &pb.RateLimitRequest{
+		Domain: "foo",
+	}).Return(nil, nil)
+	assertHttpResponse(t, handler, `{"domain": "foo"}`, 500, "text/plain; charset=utf-8", "error marshaling proto3 to json: Marshal called with nil\n")
+
+	// successful request, not rate limited
+	rls.EXPECT().ShouldRateLimit(nil, &pb.RateLimitRequest{
+		Domain: "foo",
+	}).Return(&pb.RateLimitResponse{
+		OverallCode: pb.RateLimitResponse_OK,
+	}, nil)
+	assertHttpResponse(t, handler, `{"domain": "foo"}`, 200, "application/json", `{"overallCode":"OK"}`)
+
+	// successful request, rate limited
+	rls.EXPECT().ShouldRateLimit(nil, &pb.RateLimitRequest{
+		Domain: "foo",
+	}).Return(&pb.RateLimitResponse{
+		OverallCode: pb.RateLimitResponse_OVER_LIMIT,
+	}, nil)
+	assertHttpResponse(t, handler, `{"domain": "foo"}`, 429, "application/json", `{"overallCode":"OVER_LIMIT"}`)
+}