feat: expose available needed by chatsdk

This commit is contained in:
pablo 2025-07-03 12:10:15 +03:00
parent 0ed3fc8079
commit a5ff6a1fbd
2 changed files with 38 additions and 0 deletions

View File

@ -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

View File

@ -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.