mirror of
https://github.com/logos-messaging/logos-messaging-nim.git
synced 2026-01-04 15:03:08 +00:00
feat: expose available needed by chatsdk
This commit is contained in:
parent
0ed3fc8079
commit
a5ff6a1fbd
@ -10,8 +10,13 @@
|
||||
|
||||
import testutils/unittests
|
||||
import chronos
|
||||
import std/math
|
||||
import ../../waku/common/rate_limit/token_bucket
|
||||
|
||||
# Helper function for approximate float equality
|
||||
proc equals(a, b: float, tolerance: float = 1e-6): bool =
|
||||
abs(a - b) <= tolerance
|
||||
|
||||
suite "Token Bucket":
|
||||
test "TokenBucket Sync test - strict":
|
||||
var bucket = TokenBucket.newStrict(1000, 1.milliseconds)
|
||||
@ -67,3 +72,28 @@ suite "Token Bucket":
|
||||
check bucket.tryConsume(1, start) == false
|
||||
|
||||
check bucket.tryConsume(15000, start + 1.milliseconds) == true
|
||||
|
||||
test "TokenBucket getAvailableCapacityRatio":
|
||||
var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Strict)
|
||||
var reqTime = Moment.now()
|
||||
|
||||
# Test full bucket capacity ratio
|
||||
check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0
|
||||
|
||||
# Consume some tokens and check ratio
|
||||
reqTime += 1.seconds
|
||||
check bucket.tryConsume(400, reqTime) == true
|
||||
check equals(bucket.getAvailableCapacityRatio(reqTime), 0.6) # 600/1000 = 0.6
|
||||
|
||||
# Consume more tokens
|
||||
reqTime += 1.seconds
|
||||
check bucket.tryConsume(300, reqTime) == true
|
||||
check equals(bucket.getAvailableCapacityRatio(reqTime), 0.3) # 300/1000 = 0.3
|
||||
|
||||
# Test when period has elapsed (should return 1.0)
|
||||
reqTime += 1.minutes
|
||||
check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0
|
||||
|
||||
# Test with empty bucket
|
||||
check bucket.tryConsume(1000, reqTime) == true
|
||||
check equals(bucket.getAvailableCapacityRatio(reqTime), 0.0) # 0/1000 = 0.0
|
||||
|
||||
@ -109,6 +109,14 @@ proc update(bucket: TokenBucket, currentTime: Moment) =
|
||||
else:
|
||||
updateStrict(bucket, currentTime)
|
||||
|
||||
## Returns the available capacity ratio of the bucket.
|
||||
## This is a ratio of the bucket capacity that is currently available.
|
||||
## It is a float number between 0.0 (empty) and 1.0 (full).
|
||||
proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float =
|
||||
if periodElapsed(bucket, currentTime):
|
||||
return 1.0
|
||||
return bucket.budget.float / bucket.budgetCap.float
|
||||
|
||||
proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool =
|
||||
## If `tokens` are available, consume them,
|
||||
## Otherwhise, return false.
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user