Update oauth to use a delay instead of timestamp

MichaelGHSeg · web-flow · commit b133fff466fa · 2024-10-04T15:25:03.000-04:00
* Update oauth test

* Pulling in changes and fixing test
diff --git a/segment/analytics/oauth_manager.py b/segment/analytics/oauth_manager.py
@@ -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
 
diff --git a/segment/analytics/test/test_oauth.py b/segment/analytics/test/test_oauth.py
@@ -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':
@@ -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=[]):
@@ -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()
