Skip to content

Commit

Permalink
Update oauth to use a delay instead of timestamp
Browse files Browse the repository at this point in the history
* Update oauth test

* Pulling in changes and fixing test
  • Loading branch information
MichaelGHSeg authored Oct 4, 2024
1 parent e045297 commit b133fff
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 7 deletions.
8 changes: 4 additions & 4 deletions segment/analytics/oauth_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,13 +152,13 @@ def _poller_loop(self):

elif response.status_code == 429:
self.retry_count += 1
rate_limit_reset_timestamp = None
rate_limit_reset_time = None
try:
rate_limit_reset_timestamp = int(response.headers.get("X-RateLimit-Reset"))
rate_limit_reset_time = int(response.headers.get("X-RateLimit-Reset"))
except Exception as e:
self.log.error("OAuth rate limit response did not have a valid rest time: {} | {}".format(response, e))
if rate_limit_reset_timestamp:
refresh_timer_ms = rate_limit_reset_timestamp - time.time() * 1000
if rate_limit_reset_time:
refresh_timer_ms = rate_limit_reset_time * 1000
else:
refresh_timer_ms = 5 * 1000

Expand Down
6 changes: 3 additions & 3 deletions segment/analytics/test/test_oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ def json(self):
elif kwargs['url'] == 'http://127.0.0.1:400/token':
return MockResponse({"reason": "test_reason", "json_data" : {"error":"unrecoverable", "error_description":"nah"}}, 400)
elif kwargs['url'] == 'http://127.0.0.1:429/token':
return MockResponse({"reason": "test_reason", "headers" : {"X-RateLimit-Reset": time.time()*1000 + 2000}}, 429)
return MockResponse({"reason": "test_reason", "headers" : {"X-RateLimit-Reset": 234}}, 429)
elif kwargs['url'] == 'http://127.0.0.1:500/token':
return MockResponse({"reason": "test_reason", "json_data" : {"error":"recoverable", "error_description":"nah"}}, 500)
elif kwargs['url'] == 'http://127.0.0.1:501/token':
Expand Down Expand Up @@ -106,7 +106,7 @@ def test_oauth_fail_with_retries(self, mock_post):
def test_oauth_rate_limit_delay(self, mock_sleep, mock_post):
manager = segment.analytics.oauth_manager.OauthManager("id", privatekey, "keyid", "http://127.0.0.1:429")
manager._poller_loop()
self.assertTrue(mock_sleep.call_args[0][0] > 1.9 and mock_sleep.call_args[0][0] <= 2.0)
mock_sleep.assert_called_with(234)

class TestOauthIntegration(unittest.TestCase):
def fail(self, e, batch=[]):
Expand Down Expand Up @@ -152,4 +152,4 @@ def test_oauth_integration_fail_bad_key(self, mock_post):
self.assertTrue(self.failed)

if __name__ == '__main__':
unittest.main()
unittest.main()

0 comments on commit b133fff

Please sign in to comment.