Followups to v3 upgrade (#155)

- Regenerate mocks based on new default protocol - Manually transform v2 messages to v3 messages - some of the fields were renamed thus json Marshal/Unmarshal does not work anymore - Added tests that verify conversion v2<->v3 works for headers fields - Update tests to use proto.Equal - simple assert.Equals might not work correctly for protobuf messages. Signed-off-by: Petr Pchelko <[email protected]>
envoyproxy · Jul 13, 2020 · 25e3d14 · 25e3d14
1 parent e91321b
commit 25e3d14
Show file tree

Hide file tree

Showing 10 changed files with 166 additions and 84 deletions.
diff --git a/src/service/ratelimit_legacy.go b/src/service/ratelimit_legacy.go
@@ -1,9 +1,10 @@
 package ratelimit
 
 import (
+	core_legacy "github.com/envoyproxy/go-control-plane/envoy/api/v2/core"
+	pb_struct "github.com/envoyproxy/go-control-plane/envoy/extensions/common/ratelimit/v3"
 	pb_legacy "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v2"
 	pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3"
-	"github.com/golang/protobuf/jsonpb"
 	"github.com/lyft/gostats"
 	"golang.org/x/net/context"
 )
@@ -62,38 +63,87 @@ func ConvertLegacyRequest(legacyRequest *pb_legacy.RateLimitRequest) (*pb.RateLi
 	if legacyRequest == nil {
 		return nil, nil
 	}
-
-	m := &jsonpb.Marshaler{}
-	s, err := m.MarshalToString(legacyRequest)
-	if err != nil {
-		return nil, err
+	request := &pb.RateLimitRequest{
+		Domain:     legacyRequest.GetDomain(),
+		HitsAddend: legacyRequest.GetHitsAddend(),
 	}
-
-	req := &pb.RateLimitRequest{}
-	err = jsonpb.UnmarshalString(s, req)
-	if err != nil {
-		return nil, err
+	if legacyRequest.GetDescriptors() != nil {
+		descriptors := make([]*pb_struct.RateLimitDescriptor, len(legacyRequest.GetDescriptors()))
+		for i, descriptor := range legacyRequest.GetDescriptors() {
+			if descriptor != nil {
+				descriptors[i] = &pb_struct.RateLimitDescriptor{}
+				if descriptor.GetEntries() != nil {
+					entries := make([]*pb_struct.RateLimitDescriptor_Entry, len(descriptor.GetEntries()))
+					for j, entry := range descriptor.GetEntries() {
+						if entry != nil {
+							entries[j] = &pb_struct.RateLimitDescriptor_Entry{
+								Key:   entry.GetKey(),
+								Value: entry.GetValue(),
+							}
+						}
+					}
+					descriptors[i].Entries = entries
+				}
+			}
+		}
+		request.Descriptors = descriptors
 	}
-
-	return req, nil
+	return request, nil
 }
 
 func ConvertResponse(response *pb.RateLimitResponse) (*pb_legacy.RateLimitResponse, error) {
 	if response == nil {
 		return nil, nil
 	}
 
-	m := &jsonpb.Marshaler{}
-	s, err := m.MarshalToString(response)
-	if err != nil {
-		return nil, err
+	legacyResponse := &pb_legacy.RateLimitResponse{
+		OverallCode: pb_legacy.RateLimitResponse_Code(response.GetOverallCode()),
 	}
 
-	resp := &pb_legacy.RateLimitResponse{}
-	err = jsonpb.UnmarshalString(s, resp)
-	if err != nil {
-		return nil, err
+	if response.GetStatuses() != nil {
+		statuses := make([]*pb_legacy.RateLimitResponse_DescriptorStatus, len(response.GetStatuses()))
+		for i, status := range response.GetStatuses() {
+			if status != nil {
+				statuses[i] = &pb_legacy.RateLimitResponse_DescriptorStatus{
+					Code:           pb_legacy.RateLimitResponse_Code(status.GetCode()),
+					LimitRemaining: status.GetLimitRemaining(),
+				}
+				if status.GetCurrentLimit() != nil {
+					statuses[i].CurrentLimit = &pb_legacy.RateLimitResponse_RateLimit{
+						RequestsPerUnit: status.GetCurrentLimit().GetRequestsPerUnit(),
+						Unit:            pb_legacy.RateLimitResponse_RateLimit_Unit(status.GetCurrentLimit().GetUnit()),
+					}
+				}
+			}
+		}
+		legacyResponse.Statuses = statuses
+	}
+
+	if response.GetRequestHeadersToAdd() != nil {
+		requestHeadersToAdd := make([]*core_legacy.HeaderValue, len(response.GetRequestHeadersToAdd()))
+		for i, header := range response.GetRequestHeadersToAdd() {
+			if header != nil {
+				requestHeadersToAdd[i] = &core_legacy.HeaderValue{
+					Key:   header.GetKey(),
+					Value: header.GetValue(),
+				}
+			}
+		}
+		legacyResponse.RequestHeadersToAdd = requestHeadersToAdd
+	}
+
+	if response.GetResponseHeadersToAdd() != nil {
+		responseHeadersToAdd := make([]*core_legacy.HeaderValue, len(response.GetResponseHeadersToAdd()))
+		for i, header := range response.GetResponseHeadersToAdd() {
+			if header != nil {
+				responseHeadersToAdd[i] = &core_legacy.HeaderValue{
+					Key:   header.GetKey(),
+					Value: header.GetValue(),
+				}
+			}
+		}
+		legacyResponse.Headers = responseHeadersToAdd
 	}
 
-	return resp, nil
+	return legacyResponse, nil
 }
diff --git a/test/common/common.go b/test/common/common.go
@@ -1,6 +1,9 @@
 package common
 
 import (
+	"fmt"
+	"github.com/golang/protobuf/proto"
+	"github.com/stretchr/testify/assert"
 	"sync"
 
 	pb_struct_legacy "github.com/envoyproxy/go-control-plane/envoy/api/v2/ratelimit"
@@ -69,3 +72,8 @@ func NewRateLimitRequestLegacy(domain string, descriptors [][][2]string, hitsAdd
 	request.HitsAddend = hitsAddend
 	return request
 }
+
+func AssertProtoEqual(assert *assert.Assertions, expected proto.Message, actual proto.Message) {
+	assert.True(proto.Equal(expected, actual),
+		fmt.Sprintf("These two protobuf messages are not equal:\nexpected: %v\nactual:  %v", expected, actual))
+}
diff --git a/test/integration/integration_test.go b/test/integration/integration_test.go
@@ -166,7 +166,8 @@ func testBasicBaseConfig(grpcPort, perSecond string, local_cache_size string) fu
 		response, err := c.ShouldRateLimit(
 			context.Background(),
 			common.NewRateLimitRequest("foo", [][][2]string{{{getCacheKey("hello", enable_local_cache), "world"}}}, 1))
-		assert.Equal(
+		common.AssertProtoEqual(
+			assert,
 			&pb.RateLimitResponse{
 				OverallCode: pb.RateLimitResponse_OK,
 				Statuses:    []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: nil, LimitRemaining: 0}}},
@@ -184,7 +185,8 @@ func testBasicBaseConfig(grpcPort, perSecond string, local_cache_size string) fu
 		response, err = c.ShouldRateLimit(
 			context.Background(),
 			common.NewRateLimitRequest("basic", [][][2]string{{{getCacheKey("key1", enable_local_cache), "foo"}}}, 1))
-		assert.Equal(
+		common.AssertProtoEqual(
+			assert,
 			&pb.RateLimitResponse{
 				OverallCode: pb.RateLimitResponse_OK,
 				Statuses: []*pb.RateLimitResponse_DescriptorStatus{
@@ -224,7 +226,8 @@ func testBasicBaseConfig(grpcPort, perSecond string, local_cache_size string) fu
 				limitRemaining = 0
 			}
 
-			assert.Equal(
+			common.AssertProtoEqual(
+				assert,
 				&pb.RateLimitResponse{
 					OverallCode: status,
 					Statuses: []*pb.RateLimitResponse_DescriptorStatus{
@@ -287,7 +290,8 @@ func testBasicBaseConfig(grpcPort, perSecond string, local_cache_size string) fu
 				limitRemaining2 = 0
 			}
 
-			assert.Equal(
+			common.AssertProtoEqual(
+				assert,
 				&pb.RateLimitResponse{
 					OverallCode: status,
 					Statuses: []*pb.RateLimitResponse_DescriptorStatus{
@@ -384,7 +388,8 @@ func TestBasicConfigLegacy(t *testing.T) {
 	response, err := c.ShouldRateLimit(
 		context.Background(),
 		common.NewRateLimitRequestLegacy("foo", [][][2]string{{{"hello", "world"}}}, 1))
-	assert.Equal(
+	common.AssertProtoEqual(
+		assert,
 		&pb_legacy.RateLimitResponse{
 			OverallCode: pb_legacy.RateLimitResponse_OK,
 			Statuses:    []*pb_legacy.RateLimitResponse_DescriptorStatus{{Code: pb_legacy.RateLimitResponse_OK, CurrentLimit: nil, LimitRemaining: 0}}},

diff --git a/test/mocks/config/config.go b/test/mocks/config/config.go
diff --git a/test/mocks/limiter/limiter.go b/test/mocks/limiter/limiter.go
diff --git a/test/mocks/mocks.go b/test/mocks/mocks.go
@@ -5,4 +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
+//go:generate go run github.com/golang/mock/mockgen -destination ./rls/rls.go github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3 RateLimitServiceServer
diff --git a/test/mocks/rls/rls.go b/test/mocks/rls/rls.go
diff --git a/test/server/server_impl_test.go b/test/server/server_impl_test.go
@@ -2,6 +2,8 @@ package server_test
 
 import (
 	"fmt"
+	"github.com/golang/protobuf/proto"
+	"github.com/stretchr/testify/mock"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -11,7 +13,7 @@ import (
 	pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3"
 
 	"github.com/envoyproxy/ratelimit/src/server"
-	mock_v2 "github.com/envoyproxy/ratelimit/test/mocks/rls"
+	mock_v3 "github.com/envoyproxy/ratelimit/test/mocks/rls"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 )
@@ -41,8 +43,13 @@ func TestJsonHandler(t *testing.T) {
 	controller := gomock.NewController(t)
 	defer controller.Finish()
 
-	rls := mock_v2.NewMockRateLimitServiceServer(controller)
+	rls := mock_v3.NewMockRateLimitServiceServer(controller)
 	handler := server.NewJsonHandler(rls)
+	requestMatcher := mock.MatchedBy(func(req *pb.RateLimitRequest) bool {
+		return proto.Equal(req, &pb.RateLimitRequest{
+			Domain: "foo",
+		})
+	})
 
 	// Missing request body
 	assertHttpResponse(t, handler, "", 400, "text/plain; charset=utf-8", "EOF\n")
@@ -51,35 +58,25 @@ func TestJsonHandler(t *testing.T) {
 	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)
+	rls.EXPECT().ShouldRateLimit(nil, requestMatcher).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"))
+	rls.EXPECT().ShouldRateLimit(nil, requestMatcher).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)
+	rls.EXPECT().ShouldRateLimit(nil, requestMatcher).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{
+	rls.EXPECT().ShouldRateLimit(nil, requestMatcher).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{
+	rls.EXPECT().ShouldRateLimit(nil, requestMatcher).Return(&pb.RateLimitResponse{
 		OverallCode: pb.RateLimitResponse_OVER_LIMIT,
 	}, nil)
 	assertHttpResponse(t, handler, `{"domain": "foo"}`, 429, "application/json", `{"overallCode":"OVER_LIMIT"}`)