async def test_token_bucket_get_num_tokens(): bucket = TokenBucket(1, 10) # starts at full capacity assert bucket.get_num_tokens() == 10 await bucket.take(5) assert 5 <= bucket.get_num_tokens() <= 5.1 await bucket.take(bucket.get_num_tokens()) assert 0 <= bucket.get_num_tokens() <= 0.1
async def test_token_bucket_can_take(): bucket = TokenBucket(1, 10) assert bucket.can_take() is True # can take 1 assert bucket.can_take( bucket.get_num_tokens()) is True # can take full capacity await bucket.take(10) # empty the bucket assert bucket.can_take() is False