Skip to content

Commit

Permalink
more different tests
Browse files Browse the repository at this point in the history 
Signed-off-by: wangbaiping/wbpcode <wangbaiping@bytedance.com>
  • Loading branch information
@wbpcode
wbpcode committed Dec 23, 2024
1 parent 7015ab3 commit df1198a
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 42 deletions.
10 changes: 10 additions & 0 deletions test/common/common.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import (

"github.com/stretchr/testify/assert"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/wrapperspb"

pb_struct "github.com/envoyproxy/go-control-plane/envoy/extensions/common/ratelimit/v3"
pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3"
Expand Down Expand Up @@ -64,6 +65,15 @@ func NewRateLimitRequest(domain string, descriptors [][][2]string, hitsAddend ui
return request
}

func NewRateLimitRequestWithPerDescriptorHitsAddend(domain string, descriptors [][][2]string,
hitsAddends []uint64) *pb.RateLimitRequest {
request := NewRateLimitRequest(domain, descriptors, 1)
for i, hitsAddend := range hitsAddends {
request.Descriptors[i].HitsAddend = &wrapperspb.UInt64Value{Value: 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))
Expand Down
14 changes: 7 additions & 7 deletions test/memcached/cache_impl_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,14 +102,14 @@ func TestMemcached(t *testing.T) {
nil,
)
client.EXPECT().Increment("domain_key3_value3_997200", uint64(1)).Return(uint64(11), nil)
client.EXPECT().Increment("domain_key3_value3_subkey3_subvalue3_950400", uint64(1)).Return(uint64(13), nil)
client.EXPECT().Increment("domain_key3_value3_subkey3_subvalue3_950400", uint64(2)).Return(uint64(13), nil)

request = common.NewRateLimitRequest(
request = common.NewRateLimitRequestWithPerDescriptorHitsAddend(
"domain",
[][][2]string{
{{"key3", "value3"}},
{{"key3", "value3"}, {"subkey3", "subvalue3"}},
}, 1)
}, []uint64{1, 2})
limits = []*config.RateLimit{
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_HOUR, sm.NewStats("key3_value3"), false, false, "", nil, false),
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_DAY, sm.NewStats("key3_value3_subkey3_subvalue3"), false, false, "", nil, false),
Expand All @@ -124,10 +124,10 @@ func TestMemcached(t *testing.T) {
assert.Equal(uint64(1), limits[0].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.NearLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(1), limits[0].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.NearLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(2), limits[1].Stats.TotalHits.Value())
assert.Equal(uint64(2), limits[1].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[1].Stats.NearLimit.Value())
assert.Equal(uint64(0), limits[1].Stats.WithinLimit.Value())

cache.Flush()
}
Expand Down
72 changes: 37 additions & 35 deletions test/redis/fixed_cache_impl_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,12 +86,12 @@ func testRedis(usePerSecondRedis bool) func(*testing.T) {
"EXPIRE", "domain_key2_value2_subkey2_subvalue2_1200", int64(60)).DoAndReturn(pipeAppend)
clientUsed.EXPECT().PipeDo(gomock.Any()).Return(nil)

request = common.NewRateLimitRequest(
request = common.NewRateLimitRequestWithPerDescriptorHitsAddend(
"domain",
[][][2]string{
{{"key2", "value2"}},
{{"key2", "value2"}, {"subkey2", "subvalue2"}},
}, 1)
}, []uint64{0, 1})
limits = []*config.RateLimit{
nil,
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, sm.NewStats("key2_value2_subkey2_subvalue2"), false, false, "", nil, false),
Expand All @@ -109,20 +109,20 @@ func testRedis(usePerSecondRedis bool) func(*testing.T) {

clientUsed = client
timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(5)
clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key3_value3_997200", uint64(1)).SetArg(1, uint64(11)).DoAndReturn(pipeAppend)
clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key3_value3_997200", uint64(0)).SetArg(1, uint64(11)).DoAndReturn(pipeAppend)
clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
"EXPIRE", "domain_key3_value3_997200", int64(3600)).DoAndReturn(pipeAppend)
clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key3_value3_subkey3_subvalue3_950400", uint64(1)).SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
"EXPIRE", "domain_key3_value3_subkey3_subvalue3_950400", int64(86400)).DoAndReturn(pipeAppend)
clientUsed.EXPECT().PipeDo(gomock.Any()).Return(nil)

request = common.NewRateLimitRequest(
request = common.NewRateLimitRequestWithPerDescriptorHitsAddend(
"domain",
[][][2]string{
{{"key3", "value3"}},
{{"key3", "value3"}, {"subkey3", "subvalue3"}},
}, 1)
}, []uint64{0, 1})
limits = []*config.RateLimit{
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_HOUR, sm.NewStats("key3_value3"), false, false, "", nil, false),
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_DAY, sm.NewStats("key3_value3_subkey3_subvalue3"), false, false, "", nil, false),
Expand All @@ -133,14 +133,14 @@ func testRedis(usePerSecondRedis bool) func(*testing.T) {
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0, DurationUntilReset: utils.CalculateReset(&limits[1].Limit.Unit, timeSource)},
},
cache.DoLimit(context.Background(), request, limits))
assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(1), limits[0].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.NearLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(1), limits[0].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.NearLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(1), limits[1].Stats.TotalHits.Value())
assert.Equal(uint64(1), limits[1].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[1].Stats.NearLimit.Value())
assert.Equal(uint64(0), limits[1].Stats.WithinLimit.Value())
}
}

Expand Down Expand Up @@ -644,8 +644,8 @@ func TestOverLimitWithStopCacheKeyIncrementWhenOverlimitConfig(t *testing.T) {

// Test Near Limit Stats. Under Near Limit Ratio
timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(5)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key4_value4_997200").SetArg(1, uint64(11)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key5_value5_997200").SetArg(1, uint64(11)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key4_value4_997200").SetArg(1, uint64(10)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key5_value5_997200").SetArg(1, uint64(10)).DoAndReturn(pipeAppend)

client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint64(1)).SetArg(1, uint64(11)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
Expand All @@ -656,7 +656,7 @@ func TestOverLimitWithStopCacheKeyIncrementWhenOverlimitConfig(t *testing.T) {
"EXPIRE", "domain_key5_value5_997200", int64(3600)).DoAndReturn(pipeAppend)
client.EXPECT().PipeDo(gomock.Any()).Return(nil).Times(2)

request := common.NewRateLimitRequest("domain", [][][2]string{{{"key4", "value4"}}, {{"key5", "value5"}}}, 1)
request := common.NewRateLimitRequestWithPerDescriptorHitsAddend("domain", [][][2]string{{{"key4", "value4"}}, {{"key5", "value5"}}}, []uint64{1, 1})

limits := []*config.RateLimit{
config.NewRateLimit(15, pb.RateLimitResponse_RateLimit_HOUR, sm.NewStats("key4_value4"), false, false, "", nil, false),
Expand All @@ -683,65 +683,67 @@ func TestOverLimitWithStopCacheKeyIncrementWhenOverlimitConfig(t *testing.T) {
// Check the local cache stats.
testLocalCacheStats(localCacheStats, statsStore, sink, 0, 1, 1, 0, 0)

// Test Near Limit Stats. At Near Limit Ratio, still OK
// Test Near Limit Stats. Some hits at Near Limit Ratio, but still OK.
timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(5)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key4_value4_997200").SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key5_value5_997200").SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint64(1)).SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key4_value4_997200").SetArg(1, uint64(11)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key5_value5_997200").SetArg(1, uint64(11)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint64(2)).SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
"EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key5_value5_997200", uint64(1)).SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key5_value5_997200", uint64(2)).SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
"EXPIRE", "domain_key5_value5_997200", int64(3600)).DoAndReturn(pipeAppend)
client.EXPECT().PipeDo(gomock.Any()).Return(nil).Times(2)

request = common.NewRateLimitRequestWithPerDescriptorHitsAddend("domain", [][][2]string{{{"key4", "value4"}}, {{"key5", "value5"}}}, []uint64{2, 2})

assert.Equal(
[]*pb.RateLimitResponse_DescriptorStatus{
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 2, DurationUntilReset: utils.CalculateReset(&limits[0].Limit.Unit, timeSource)},
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 1, DurationUntilReset: utils.CalculateReset(&limits[1].Limit.Unit, timeSource)},
},
cache.DoLimit(context.Background(), request, limits))
assert.Equal(uint64(2), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.OverLimitWithLocalCache.Value())
assert.Equal(uint64(1), limits[0].Stats.NearLimit.Value())
assert.Equal(uint64(2), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(2), limits[1].Stats.TotalHits.Value())
assert.Equal(uint64(3), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(3), limits[1].Stats.TotalHits.Value())
assert.Equal(uint64(0), limits[1].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[1].Stats.OverLimitWithLocalCache.Value())
assert.Equal(uint64(1), limits[1].Stats.NearLimit.Value())
assert.Equal(uint64(2), limits[1].Stats.WithinLimit.Value())
assert.Equal(uint64(2), limits[1].Stats.NearLimit.Value())
assert.Equal(uint64(3), limits[1].Stats.WithinLimit.Value())

// Check the local cache stats.
testLocalCacheStats(localCacheStats, statsStore, sink, 0, 2, 2, 0, 0)

// Test one key is reaching to the Overlimit threshold
timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(5)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key4_value4_997200").SetArg(1, uint64(14)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key5_value5_997200").SetArg(1, uint64(14)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint64(0)).SetArg(1, uint64(14)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key4_value4_997200").SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "GET", "domain_key5_value5_997200").SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint64(0)).SetArg(1, uint64(13)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
"EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key5_value5_997200", uint64(1)).SetArg(1, uint64(14)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key5_value5_997200", uint64(2)).SetArg(1, uint64(15)).DoAndReturn(pipeAppend)
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
"EXPIRE", "domain_key5_value5_997200", int64(3600)).DoAndReturn(pipeAppend)
client.EXPECT().PipeDo(gomock.Any()).Return(nil).Times(2)

assert.Equal(
[]*pb.RateLimitResponse_DescriptorStatus{
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 1, DurationUntilReset: utils.CalculateReset(&limits[0].Limit.Unit, timeSource)},
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0, DurationUntilReset: utils.CalculateReset(&limits[1].Limit.Unit, timeSource)},
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 2, DurationUntilReset: utils.CalculateReset(&limits[0].Limit.Unit, timeSource)},
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0, DurationUntilReset: utils.CalculateReset(&limits[1].Limit.Unit, timeSource)},
},
cache.DoLimit(context.Background(), request, limits))
assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(5), limits[0].Stats.TotalHits.Value())
assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[0].Stats.OverLimitWithLocalCache.Value())
assert.Equal(uint64(2), limits[0].Stats.NearLimit.Value())
assert.Equal(uint64(3), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(3), limits[1].Stats.TotalHits.Value())
assert.Equal(uint64(0), limits[1].Stats.OverLimit.Value())
assert.Equal(uint64(5), limits[0].Stats.WithinLimit.Value())
assert.Equal(uint64(5), limits[1].Stats.TotalHits.Value())
assert.Equal(uint64(1), limits[1].Stats.OverLimit.Value())
assert.Equal(uint64(0), limits[1].Stats.OverLimitWithLocalCache.Value())
assert.Equal(uint64(2), limits[1].Stats.NearLimit.Value())
assert.Equal(uint64(3), limits[1].Stats.NearLimit.Value())
assert.Equal(uint64(3), limits[1].Stats.WithinLimit.Value())

// Check the local cache stats.
Expand Down

0 comments on commit df1198a

Please sign in to comment.