async def test_start(self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
assert mgr.gc_task is None
mgr.start()
mgr.start()
mgr.start()
assert mgr.gc_task is not None
async def test_start(self):
with buckets.RESTBucketManager() as mgr:
assert mgr.gc_task is None
mgr.start()
mgr.start()
mgr.start()
assert mgr.gc_task is not None
async def test_exit_closes(self):
with mock.patch.object(buckets.RESTBucketManager, "close") as close:
with mock.patch.object(buckets.RESTBucketManager, "gc") as gc:
with buckets.RESTBucketManager() as mgr:
mgr.start(0.01, 32)
gc.assert_called_once_with(0.01, 32)
close.assert_called()
async def test_acquire_route_when_not_in_routes_to_real_hashes_doesnt_cache_route(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda initial_hash: initial_hash + ";bobs")
mgr.acquire(route)
assert mgr.routes_to_hashes.get(route.route) is None
async def test_acquire_route_returns_acquired_future_for_new_bucket(self):
with buckets.RESTBucketManager() as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
return_value="eat pant;bobs")
bucket = mock.Mock()
mgr.routes_to_hashes[route.route] = "eat pant"
mgr.real_hashes_to_buckets["eat pant;bobs"] = bucket
f = mgr.acquire(route)
assert f is bucket.acquire()
async def test_acquire_route_returns_acquired_future(self):
with buckets.RESTBucketManager() as mgr:
route = mock.Mock()
bucket = mock.Mock()
with mock.patch.object(buckets, "RESTBucket", return_value=bucket):
route.create_real_bucket_hash = mock.Mock(
wraps=lambda intial_hash: intial_hash + ";bobs")
f = mgr.acquire(route)
assert f is bucket.acquire()
async def test_acquire_route_when_not_in_routes_to_real_hashes_caches_route(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda intial_hash: intial_hash + ";bobs")
# This isn't a coroutine; why would I await it?
mgr.acquire(route)
assert mgr.routes_to_hashes[route.route] == "UNKNOWN"
async def test_acquire_route_returns_context_manager(self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
bucket = mock.Mock(reset_at=time.perf_counter() +
999999999999999999999999999)
with mock.patch.object(buckets, "RESTBucket", return_value=bucket):
route.create_real_bucket_hash = mock.Mock(
wraps=lambda initial_hash: initial_hash + ";bobs")
assert mgr.acquire(route) is bucket
async def test_update_rate_limits_if_wrong_bucket_hash_reroutes_route(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda intial_hash: intial_hash + ";bobs")
mgr.routes_to_hashes[route.route] = "123"
mgr.update_rate_limits(route, "blep", 22, 23, 3.56)
assert mgr.routes_to_hashes[route.route] == "blep"
assert isinstance(mgr.real_hashes_to_buckets["blep;bobs"],
buckets.RESTBucket)
async def test_acquire_route_when_route_cached_already_obtains_hash_from_route_and_bucket_from_hash(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
return_value="eat pant;1234")
bucket = mock.Mock(reset_at=time.perf_counter() +
999999999999999999999999999)
mgr.routes_to_hashes[route.route] = "eat pant"
mgr.real_hashes_to_buckets["eat pant;1234"] = bucket
assert mgr.acquire(route) is bucket
async def test_acquire_unknown_route_returns_context_manager_for_new_bucket(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
return_value="eat pant;bobs")
bucket = mock.Mock(reset_at=time.perf_counter() +
999999999999999999999999999)
mgr.routes_to_hashes[route.route] = "eat pant"
mgr.real_hashes_to_buckets["eat pant;bobs"] = bucket
assert mgr.acquire(route) is bucket
async def test_acquire_route_when_too_long_ratelimit(self):
with buckets.RESTBucketManager(max_rate_limit=60) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
return_value="eat pant;bobs")
bucket = mock.Mock(reset_at=time.perf_counter() +
999999999999999999999999999,
is_ratelimited=mock.Mock(return_value=True))
mgr.routes_to_hashes[route.route] = "eat pant"
mgr.real_hashes_to_buckets["eat pant;bobs"] = bucket
with pytest.raises(errors.RateLimitTooLongError):
mgr.acquire(route)
async def test_acquire_route_when_not_in_routes_to_real_hashes_makes_new_bucket_using_initial_hash(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda intial_hash: intial_hash + ";bobs")
# This isn't a coroutine; why would I await it?
mgr.acquire(route)
assert "UNKNOWN;bobs" in mgr.real_hashes_to_buckets
assert isinstance(mgr.real_hashes_to_buckets["UNKNOWN;bobs"],
buckets.RESTBucket)
async def test_update_rate_limits_if_right_bucket_hash_does_nothing_to_hash(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda intial_hash: intial_hash + ";bobs")
mgr.routes_to_hashes[route.route] = "123"
bucket = mock.Mock(reset_at=time.perf_counter() +
999999999999999999999999999)
mgr.real_hashes_to_buckets["123;bobs"] = bucket
mgr.update_rate_limits(route, "123", 22, 23, 7.65)
assert mgr.routes_to_hashes[route.route] == "123"
assert mgr.real_hashes_to_buckets["123;bobs"] is bucket
async def test_update_rate_limits_if_right_bucket_hash_does_nothing_to_hash(
self):
with buckets.RESTBucketManager() as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda intial_hash: intial_hash + ";bobs")
mgr.routes_to_hashes[route.route] = "123"
bucket = mock.Mock()
mgr.real_hashes_to_buckets["123;bobs"] = bucket
mgr.update_rate_limits(route, "123", 22, 23,
datetime.datetime.now(),
datetime.datetime.now())
assert mgr.routes_to_hashes[route.route] == "123"
assert mgr.real_hashes_to_buckets["123;bobs"] is bucket
async def test_acquire_route_when_not_in_routes_to_real_hashes_makes_new_bucket_using_initial_hash(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
with mock.patch.object(
buckets, "_create_unknown_hash",
return_value="UNKNOWN;bobs") as create_unknown_hash:
mgr.acquire(route)
assert "UNKNOWN;bobs" in mgr.real_hashes_to_buckets
assert isinstance(mgr.real_hashes_to_buckets["UNKNOWN;bobs"],
buckets.RESTBucket)
create_unknown_hash.assert_called_once_with(route)
async def test_update_rate_limits_updates_params(self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda initial_hash: initial_hash + ";bobs")
mgr.routes_to_hashes[route.route] = "123"
bucket = mock.Mock(reset_at=time.perf_counter() +
999999999999999999999999999)
mgr.real_hashes_to_buckets["123;bobs"] = bucket
with mock.patch.object(hikari_date, "monotonic", return_value=27):
mgr.update_rate_limits(route, "123", 22, 23, 5.32)
bucket.update_rate_limit.assert_called_once_with(
22, 23, 27 + 5.32)
async def test_acquire_route_when_route_cached_already_obtains_hash_from_route_and_bucket_from_hash(
self):
with buckets.RESTBucketManager() as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
return_value="eat pant;1234")
bucket = mock.Mock()
mgr.routes_to_hashes[route] = "eat pant"
mgr.real_hashes_to_buckets["eat pant;1234"] = bucket
# This isn't a coroutine; why would I await it?
mgr.acquire(route)
# yes i test this twice, sort of. no, there isn't another way to verify this. sue me.
bucket.acquire.assert_called_once()
async def test_update_rate_limits_if_wrong_bucket_hash_reroutes_route(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda initial_hash: initial_hash + ";bobs")
mgr.routes_to_hashes[route.route] = "123"
with mock.patch.object(hikari_date, "monotonic", return_value=27):
with mock.patch.object(buckets, "RESTBucket") as bucket:
mgr.update_rate_limits(route, "blep", 22, 23, 3.56)
assert mgr.routes_to_hashes[route.route] == "blep"
assert mgr.real_hashes_to_buckets[
"blep;bobs"] is bucket.return_value
bucket.return_value.update_rate_limit.assert_called_once_with(
22, 23, 27 + 3.56)
async def test_gc_polls_until_closed_event_set(self):
# This is shit, but it is good shit.
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
mgr.start(0.01)
assert mgr.gc_task is not None
assert not mgr.gc_task.done()
await hikari_test_helpers.idle()
assert mgr.gc_task is not None
assert not mgr.gc_task.done()
await hikari_test_helpers.idle()
mgr.closed_event.set()
assert mgr.gc_task is not None
assert not mgr.gc_task.done()
task = mgr.gc_task
await hikari_test_helpers.idle()
assert mgr.gc_task is None
assert task.done()
async def test_update_rate_limits_updates_params(self):
with buckets.RESTBucketManager() as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda intial_hash: intial_hash + ";bobs")
mgr.routes_to_hashes[route.route] = "123"
bucket = mock.Mock()
mgr.real_hashes_to_buckets["123;bobs"] = bucket
date = datetime.datetime.now().replace(year=2004)
reset_at = datetime.datetime.now()
with mock.patch.object(hikari_date, "monotonic", return_value=27):
expect_reset_at_monotonic = 27 + (reset_at -
date).total_seconds()
mgr.update_rate_limits(route, "123", 22, 23, date, reset_at)
bucket.update_rate_limit.assert_called_once_with(
22, 23, expect_reset_at_monotonic)
async def test_close_closes_all_buckets(self):
class MockBucket:
def __init__(self):
self.close = mock.Mock()
buckets_array = [MockBucket() for _ in range(30)]
mgr = buckets.RESTBucketManager(max_rate_limit=float("inf"))
mgr.real_hashes_to_buckets = {
f"blah{i}": bucket
for i, bucket in enumerate(buckets_array)
}
mgr.close()
for i, bucket in enumerate(buckets_array):
bucket.close.assert_called_once(), i
async def test_gc_polls_until_closed_event_set(self, event_loop):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
# Start the gc and initial assertions
task = event_loop.create_task(mgr.gc(0.001, float("inf")))
assert not task.done()
# [First poll] event not set => shouldn't complete the task
await asyncio.sleep(0.001)
assert not task.done()
# [Second poll] event not set during poll => shouldn't complete the task
await asyncio.sleep(0.001)
mgr.closed_event.set()
assert not task.done()
# [Third poll] event set => should complete the task
await asyncio.sleep(0.001)
assert task.done()
async def test_update_rate_limits_if_unknown_bucket_hash_reroutes_route(
self):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
route = mock.Mock()
route.create_real_bucket_hash = mock.Mock(
wraps=lambda initial_hash: initial_hash + ";bobs")
mgr.routes_to_hashes[route.route] = "123"
bucket = mock.Mock()
mgr.real_hashes_to_buckets["UNKNOWN;bobs"] = bucket
with mock.patch.object(
buckets, "_create_unknown_hash",
return_value="UNKNOWN;bobs") as create_unknown_hash:
with mock.patch.object(hikari_date,
"monotonic",
return_value=27):
mgr.update_rate_limits(route, "blep", 22, 23, 3.56)
assert mgr.routes_to_hashes[route.route] == "blep"
assert mgr.real_hashes_to_buckets["blep;bobs"] is bucket
bucket.resolve.assert_called_once_with("blep;bobs")
bucket.update_rate_limit.assert_called_once_with(22, 23, 27 + 3.56)
create_unknown_hash.assert_called_once_with(route)
async def test_close_sets_closed_event(self):
mgr = buckets.RESTBucketManager(max_rate_limit=float("inf"))
assert not mgr.closed_event.is_set()
mgr.close()
assert mgr.closed_event.is_set()
def test_is_started(self, gc_task, is_started):
with buckets.RESTBucketManager() as mgr:
mgr.gc_task = gc_task
assert mgr.is_started is is_started
def test_is_started(self, gc_task, is_started):
with buckets.RESTBucketManager(max_rate_limit=float("inf")) as mgr:
mgr.gc_task = gc_task
assert mgr.is_started is is_started
async def test_close_sets_closed_event(self):
mgr = buckets.RESTBucketManager()
assert not mgr.closed_event.is_set()
mgr.close()
assert mgr.closed_event.is_set()