@@ -20,7 +20,7 @@ def emit(self, record):
20
20
"""Check record to see if it is a rate-limit message."""
21
21
assert isinstance (record , logging .LogRecord )
22
22
23
- if record . msg . startswith ( "Received rate-limit message" ) :
23
+ if "Received rate-limit message" in record . msg :
24
24
self .rate_limit_detected = True
25
25
26
26
@@ -29,17 +29,19 @@ class TestRestSession:
29
29
"""Test edge cases of core RestSession functionality."""
30
30
31
31
@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 ):
33
33
logger = logging .getLogger (__name__ )
34
34
35
35
# 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
38
38
39
39
# Add log handler
40
40
root_logger = logging .getLogger ()
41
41
rate_limit_detector = RateLimitDetector ()
42
42
root_logger .addHandler (rate_limit_detector )
43
+ logger = logging .getLogger (__name__ )
44
+ logger .info ("Starting Rate Limit Testing" )
43
45
44
46
try :
45
47
# Try and trigger a rate-limit
@@ -56,7 +58,7 @@ def test_rate_limit_retry_indefinitely(self, api, rooms_list, add_rooms):
56
58
request_count )
57
59
# Remove the log handler and restore the pre-test state
58
60
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
60
62
61
63
# Assert test condition
62
64
assert rate_limit_detector .rate_limit_detected == True
0 commit comments