logos-messaging-nim/tests/common/test_tokenbucket.nim
2025-07-03 16:58:14 +03:00

100 lines
3.4 KiB
Nim

# Chronos Test Suite
# (c) Copyright 2022-Present
# Status Research & Development GmbH
#
# Licensed under either of
# Apache License, version 2.0, (LICENSE-APACHEv2)
# MIT license (LICENSE-MIT)
{.used.}
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)
let
start = Moment.now()
fullTime = start + 1.milliseconds
check:
bucket.tryConsume(800, start) == true
bucket.tryConsume(200, start) == true
# Out of budget
bucket.tryConsume(100, start) == false
bucket.tryConsume(800, fullTime) == true
bucket.tryConsume(200, fullTime) == true
# Out of budget
bucket.tryConsume(100, fullTime) == false
test "TokenBucket Sync test - compensating":
var bucket = TokenBucket.new(1000, 1.milliseconds)
let
start = Moment.now()
fullTime = start + 1.milliseconds
check:
bucket.tryConsume(800, start) == true
bucket.tryConsume(200, start) == true
# Out of budget
bucket.tryConsume(100, start) == false
bucket.tryConsume(800, fullTime) == true
bucket.tryConsume(200, fullTime) == true
# Due not using the bucket for a full period the compensation will satisfy this request
bucket.tryConsume(100, fullTime) == true
test "TokenBucket Max compensation":
var bucket = TokenBucket.new(1000, 1.minutes)
var reqTime = Moment.now()
check bucket.tryConsume(1000, reqTime)
check bucket.tryConsume(1, reqTime) == false
reqTime += 1.minutes
check bucket.tryConsume(500, reqTime) == true
reqTime += 1.minutes
check bucket.tryConsume(1000, reqTime) == true
reqTime += 10.seconds
# max compensation is 25% so try to consume 250 more
check bucket.tryConsume(250, reqTime) == true
reqTime += 49.seconds
# out of budget within the same period
check bucket.tryConsume(1, reqTime) == false
test "TokenBucket Short replenish":
var bucket = TokenBucket.new(15000, 1.milliseconds)
let start = Moment.now()
check bucket.tryConsume(15000, start)
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