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