1// Copyright (c) 2017 Uber Technologies, Inc. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package utils 16 17import ( 18 "testing" 19 "time" 20 21 "github.com/stretchr/testify/assert" 22) 23 24func TestRateLimiter(t *testing.T) { 25 rl := NewRateLimiter(2.0, 2.0) 26 // stop time 27 ts := time.Now() 28 rl.lastTick = ts 29 rl.timeNow = func() time.Time { 30 return ts 31 } 32 assert.True(t, rl.CheckCredit(1.0)) 33 assert.True(t, rl.CheckCredit(1.0)) 34 assert.False(t, rl.CheckCredit(1.0)) 35 // move time 250ms forward, not enough credits to pay for 1.0 item 36 rl.timeNow = func() time.Time { 37 return ts.Add(time.Second / 4) 38 } 39 assert.False(t, rl.CheckCredit(1.0)) 40 // move time 500ms forward, now enough credits to pay for 1.0 item 41 rl.timeNow = func() time.Time { 42 return ts.Add(time.Second/4 + time.Second/2) 43 } 44 assert.True(t, rl.CheckCredit(1.0)) 45 assert.False(t, rl.CheckCredit(1.0)) 46 // move time 5s forward, enough to accumulate credits for 10 messages, but it should still be capped at 2 47 rl.lastTick = ts 48 rl.timeNow = func() time.Time { 49 return ts.Add(5 * time.Second) 50 } 51 assert.True(t, rl.CheckCredit(1.0)) 52 assert.True(t, rl.CheckCredit(1.0)) 53 assert.False(t, rl.CheckCredit(1.0)) 54 assert.False(t, rl.CheckCredit(1.0)) 55 assert.False(t, rl.CheckCredit(1.0)) 56} 57 58func TestRateLimiterMaxBalance(t *testing.T) { 59 rl := NewRateLimiter(0.1, 1.0) 60 // stop time 61 ts := time.Now() 62 rl.lastTick = ts 63 rl.timeNow = func() time.Time { 64 return ts 65 } 66 assert.True(t, rl.CheckCredit(1.0), "on initialization, should have enough credits for 1 message") 67 68 // move time 20s forward, enough to accumulate credits for 2 messages, but it should still be capped at 1 69 rl.timeNow = func() time.Time { 70 return ts.Add(time.Second * 20) 71 } 72 assert.True(t, rl.CheckCredit(1.0)) 73 assert.False(t, rl.CheckCredit(1.0)) 74} 75 76func TestRateLimiterReconfigure(t *testing.T) { 77 rl := NewRateLimiter(1, 1.0) 78 assertBalance := func(expected float64) { 79 const delta = 0.0000001 // just some precision for comparing floats 80 assert.InDelta(t, expected, rl.balance, delta) 81 } 82 // stop time 83 ts := time.Now() 84 rl.lastTick = ts 85 rl.timeNow = func() time.Time { 86 return ts 87 } 88 assert.True(t, rl.CheckCredit(1.0), "first message must succeed") 89 assert.False(t, rl.CheckCredit(1.0), "second message must be rejected") 90 assertBalance(0.0) 91 92 // move half-second forward 93 rl.timeNow = func() time.Time { 94 return ts.Add(time.Second / 2) 95 } 96 rl.updateBalance() 97 assertBalance(0.5) // 50% of max 98 99 rl.Update(2, 4) 100 assertBalance(2) // 50% of max 101 assert.EqualValues(t, 2, rl.creditsPerSecond) 102 assert.EqualValues(t, 4, rl.maxBalance) 103} 104