52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
from backup.util import TokenBucket
|
|
from ..faketime import FakeTime
|
|
|
|
|
|
async def test_consume(time: FakeTime):
|
|
bucket = TokenBucket(time, 10, 1, 1)
|
|
assert bucket.consume(1)
|
|
assert not bucket.consume(1)
|
|
|
|
time.advance(seconds=1)
|
|
assert bucket.consume(1)
|
|
assert not bucket.consume(1)
|
|
|
|
|
|
async def test_async_consume(time: FakeTime):
|
|
bucket = TokenBucket(time, 10, 1, 1)
|
|
assert await bucket.consumeWithWait(1, 2) == 1
|
|
assert len(time.sleeps) == 0
|
|
|
|
time.advance(seconds=2)
|
|
assert await bucket.consumeWithWait(1, 2) == 2
|
|
assert len(time.sleeps) == 0
|
|
|
|
assert await bucket.consumeWithWait(1, 2) == 1
|
|
assert len(time.sleeps) == 1
|
|
assert time.sleeps[0] == 1
|
|
|
|
|
|
async def test_capacity(time: FakeTime):
|
|
bucket = TokenBucket(time, 10, 1)
|
|
assert await bucket.consumeWithWait(1, 10) == 10
|
|
assert len(time.sleeps) == 0
|
|
|
|
assert await bucket.consumeWithWait(5, 10) == 5
|
|
assert len(time.sleeps) == 1
|
|
assert time.sleeps[0] == 5
|
|
|
|
time.clearSleeps()
|
|
assert await bucket.consumeWithWait(20, 20) == 20
|
|
assert len(time.sleeps) == 1
|
|
assert time.sleeps[0] == 20
|
|
|
|
time.clearSleeps()
|
|
time.advance(seconds=5)
|
|
assert await bucket.consumeWithWait(1, 10) == 5
|
|
|
|
|
|
async def test_higher_fill_rate(time: FakeTime):
|
|
bucket = TokenBucket(time, capacity=1000, fill_rate=100)
|
|
assert await bucket.consumeWithWait(1, 1000) == 1000
|
|
assert len(time.sleeps) == 0
|