Skip to content

Commit 66205dc

Browse files
committed
Update Rate Limit Test
1 parent 6f6f32e commit 66205dc

File tree

1 file changed

+7
-5
lines changed

1 file changed

+7
-5
lines changed

tests/test_restsession.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ def emit(self, record):
2020
"""Check record to see if it is a rate-limit message."""
2121
assert isinstance(record, logging.LogRecord)
2222

23-
if record.msg.startswith("Received rate-limit message"):
23+
if "Received rate-limit message" in record.msg:
2424
self.rate_limit_detected = True
2525

2626

@@ -29,17 +29,19 @@ class TestRestSession:
2929
"""Test edge cases of core RestSession functionality."""
3030

3131
@pytest.mark.ratelimit
32-
def test_rate_limit_retry_indefinitely(self, api, rooms_list, add_rooms):
32+
def test_rate_limit_retry(self, api, rooms_list, add_rooms):
3333
logger = logging.getLogger(__name__)
3434

3535
# Save state and initialize test setup
36-
original_rate_limit_timer = api._session.rate_limit_timeout
37-
api._session.rate_limit_timeout = None
36+
original_wait_on_rate_limit = api._session.wait_on_rate_limit
37+
api._session.wait_on_rate_limit = True
3838

3939
# Add log handler
4040
root_logger = logging.getLogger()
4141
rate_limit_detector = RateLimitDetector()
4242
root_logger.addHandler(rate_limit_detector)
43+
logger = logging.getLogger(__name__)
44+
logger.info("Starting Rate Limit Testing")
4345

4446
try:
4547
# Try and trigger a rate-limit
@@ -56,7 +58,7 @@ def test_rate_limit_retry_indefinitely(self, api, rooms_list, add_rooms):
5658
request_count)
5759
# Remove the log handler and restore the pre-test state
5860
root_logger.removeHandler(rate_limit_detector)
59-
api._session.rate_limit_timeout = original_rate_limit_timer
61+
api._session.wait_on_rate_limit = original_wait_on_rate_limit
6062

6163
# Assert test condition
6264
assert rate_limit_detector.rate_limit_detected == True

0 commit comments

Comments
 (0)