1+ #include < gtest/gtest.h>
2+ #include " algorithms/token_bucket.h"
3+ #include " time/timestamp.h"
4+ #include < thread>
5+ #include < chrono>
6+
7+ using namespace BaseKit ;
8+
9+ // 测试TokenBucket基本功能
10+ TEST (TokenBucketTest, BasicFunctionality) {
11+ // 创建一个每秒10个令牌，最大容量为20个令牌的桶
12+ TokenBucket bucket (10 , 20 );
13+
14+ // 测试消费单个令牌
15+ EXPECT_TRUE (bucket.Consume (1 ));
16+
17+ // 测试消费多个令牌
18+ EXPECT_TRUE (bucket.Consume (5 ));
19+
20+ // 测试拷贝构造函数
21+ TokenBucket bucket2 (bucket);
22+ EXPECT_TRUE (bucket2.Consume (1 ));
23+
24+ // 测试赋值操作符
25+ TokenBucket bucket3 (100 , 100 );
26+ bucket3 = bucket;
27+ EXPECT_TRUE (bucket3.Consume (1 ));
28+ }
29+
30+ // 测试TokenBucket速率限制功能
31+ TEST (TokenBucketTest, RateLimit) {
32+ // 创建一个每秒10个令牌，最大容量为10个令牌的桶
33+ TokenBucket bucket (10 , 10 );
34+
35+ // 立即消费所有令牌（最大容量）
36+ EXPECT_TRUE (bucket.Consume (10 ));
37+
38+ // 由于令牌已耗尽，此时无法继续消费
39+ EXPECT_FALSE (bucket.Consume (1 ));
40+
41+ // 等待一小段时间，让桶重新积累一些令牌
42+ std::this_thread::sleep_for (std::chrono::milliseconds (200 )); // 等待200毫秒，应该积累约2个令牌
43+
44+ // 尝试消费1个令牌，应该成功
45+ EXPECT_TRUE (bucket.Consume (1 ));
46+
47+ // 尝试消费2个令牌，应该失败（因为只积累了约1个令牌）
48+ EXPECT_FALSE (bucket.Consume (2 ));
49+
50+ // 再等待一段时间
51+ std::this_thread::sleep_for (std::chrono::milliseconds (300 )); // 再等待300毫秒，应该再积累约3个令牌
52+
53+ // 尝试消费3个令牌，应该成功
54+ EXPECT_TRUE (bucket.Consume (3 ));
55+ }
56+
57+ // 测试TokenBucket突发流量控制
58+ TEST (TokenBucketTest, BurstControl) {
59+ // 创建一个每秒20个令牌，最大容量为100个令牌的桶（允许短时间内的突发流量）
60+ TokenBucket bucket (20 , 100 );
61+
62+ // 初始桶是满的，不需要时间积累
63+ // 在初始状态下应该有100个令牌
64+ EXPECT_TRUE (bucket.Consume (100 ));
65+
66+ // 尝试立即消费更多令牌，应该失败
67+ EXPECT_FALSE (bucket.Consume (1 ));
68+
69+ // 等待足够长的时间，让桶积累一些令牌
70+ std::this_thread::sleep_for (std::chrono::milliseconds (500 )); // 等待500毫秒，应该积累约10个令牌
71+
72+ // 测试积累的令牌数
73+ EXPECT_TRUE (bucket.Consume (10 )); // 应该成功消费10个令牌
74+ EXPECT_FALSE (bucket.Consume (1 )); // 再消费应该失败
75+
76+ // 再等待一段时间，让桶积累更多令牌
77+ std::this_thread::sleep_for (std::chrono::seconds (1 )); // 等待1秒，应该积累约20个令牌
78+
79+ // 测试积累的令牌数
80+ EXPECT_TRUE (bucket.Consume (20 )); // 应该成功消费20个令牌
81+ EXPECT_FALSE (bucket.Consume (1 )); // 再消费应该失败
82+ }
83+
84+ // 测试TokenBucket的长时间不使用后的行为
85+ TEST (TokenBucketTest, LongInactivity) {
86+ // 创建一个每秒5个令牌，最大容量为10个令牌的桶
87+ TokenBucket bucket (5 , 10 );
88+
89+ // 等待足够长的时间，让桶积累满令牌
90+ std::this_thread::sleep_for (std::chrono::seconds (3 )); // 等待3秒，应该积累约10个令牌（最大容量）
91+
92+ // 测试桶是否已经填满
93+ EXPECT_TRUE (bucket.Consume (10 )); // 应该能消费10个令牌
94+ EXPECT_FALSE (bucket.Consume (1 )); // 桶应该为空
95+
96+ // 测试长时间不使用后再次使用
97+ std::this_thread::sleep_for (std::chrono::seconds (10 )); // 等待10秒
98+
99+ // 测试是否能够消费最大容量的令牌（桶应该已重新填满）
100+ EXPECT_TRUE (bucket.Consume (10 ));
101+ }
102+
103+ // 测试极限情况
104+ TEST (TokenBucketTest, EdgeCases) {
105+ // 测试极低速率的桶
106+ TokenBucket lowRate (1 , 1 ); // 每秒只有1个令牌
107+ EXPECT_TRUE (lowRate.Consume (1 )); // 初始应该有1个令牌
108+ EXPECT_FALSE (lowRate.Consume (1 )); // 现在应该没有令牌了
109+
110+ // 测试极高速率的桶
111+ TokenBucket highRate (1000000 , 1000 ); // 每秒100万个令牌
112+ EXPECT_TRUE (highRate.Consume (1000 )); // 应该能立即消费1000个令牌
113+
114+ // 测试0令牌消费
115+ TokenBucket normal (10 , 10 );
116+ EXPECT_TRUE (normal.Consume (0 )); // 消费0个令牌应该总是成功
117+ }
0 commit comments