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
