Support for Fitbit API rate limiting headers

fitbit/api.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 import datetime
 import json
+import time
+
 import requests
 
 try:
@@ -243,6 +245,9 @@ def __init__(self, client_id, client_secret, access_token=None,
             setattr(self, '%s_activities' % qualifier, curry(self.activity_stats, qualifier=qualifier))
             setattr(self, '%s_foods' % qualifier, curry(self._food_stats,
                                                         qualifier=qualifier))
+        self.rate_limit_remaining = None
+        self.rate_limit_reset = None
+        self.rate_limit_limit = None
 
     def make_request(self, *args, **kwargs):
         # This should handle data level errors, improper requests, and bad
@@ -252,8 +257,22 @@ def make_request(self, *args, **kwargs):
         kwargs['headers'] = headers
 
         method = kwargs.get('method', 'POST' if 'data' in kwargs else 'GET')
+
+        if (self.rate_limit_remaining is not None and
+                self.rate_limit_remaining == 0 and
+                time.time() < self.rate_limit_reset):
+            raise exceptions.RateLimited(self.rate_limit_limit, self.rate_limit_remaining, self.rate_limit_reset)
+
         response = self.client.make_request(*args, **kwargs)
 
+        if 'fitbit-rate-limit-remaining' in response.headers:
+            self.rate_limit_remaining = int(response.headers.get('fitbit-rate-limit-remaining'))
+        if 'fitbit-rate-limit-limit' in response.headers:
+            self.rate_limit_limit = int(response.headers.get('fitbit-rate-limit-limit'))
+        rate_limit_reset = response.headers.get('fitbit-rate-limit-reset')
+        if rate_limit_reset:
+            self.rate_limit_reset = time.time() + int(rate_limit_reset)
+
         if response.status_code == 202:
             return True
         if method == 'DELETE':

fitbit/exceptions.py

@@ -1,4 +1,5 @@
 import json
+import time
 
 
 class BadResponse(Exception):
@@ -22,6 +23,20 @@ class Timeout(Exception):
     pass
 
 
+class RateLimited(Exception):
+    """
+    Used when the Fitbit API rate limit has been exceeded and a request would cause an HTTP 429 error.
+    """
+
+    def __init__(self, rate_limit_limit, rate_limit_remaining, rate_limit_reset):
+        self.rate_limit_limit = rate_limit_limit
+        self.rate_limit_remaining = rate_limit_remaining
+        self.rate_limit_reset = rate_limit_reset
+        super(RateLimited, self).__init__(
+            "Rate limit of {} requests exhausted. Reset in {:0f} seconds".format(rate_limit_limit,
+                                                                                 rate_limit_reset - time.time()))
+
+
 class HTTPException(Exception):
     def __init__(self, response, *args, **kwargs):
         try:

fitbit_tests/__init__.py

@@ -3,6 +3,7 @@
 from .test_auth import Auth2Test
 from .test_api import (
     APITest,
+    RateLimitTest,
     CollectionResourceTest,
     DeleteCollectionResourceTest,
     ResourceAccessTest,
@@ -16,6 +17,7 @@ def all_tests(consumer_key="", consumer_secret="", user_key=None, user_secret=No
     suite.addTest(unittest.makeSuite(ExceptionTest))
     suite.addTest(unittest.makeSuite(Auth2Test))
     suite.addTest(unittest.makeSuite(APITest))
+    suite.addTest(unittest.makeSuite(RateLimitTest))
     suite.addTest(unittest.makeSuite(CollectionResourceTest))
     suite.addTest(unittest.makeSuite(DeleteCollectionResourceTest))
     suite.addTest(unittest.makeSuite(ResourceAccessTest))

fitbit_tests/test_api.py

@@ -1,9 +1,10 @@
+import time
 from unittest import TestCase
 import datetime
 import mock
 import requests
 from fitbit import Fitbit
-from fitbit.exceptions import DeleteError, Timeout
+from fitbit.exceptions import DeleteError, Timeout, RateLimited
 
 URLBASE = "%s/%s/user" % (Fitbit.API_ENDPOINT, Fitbit.API_VERSION)
 
@@ -82,6 +83,7 @@ def test_make_request(self):
         mock_response = mock.Mock()
         mock_response.status_code = 200
         mock_response.content = b"1"
+        mock_response.headers = {}
         with mock.patch.object(self.fb.client, 'make_request') as client_make_request:
             client_make_request.return_value = mock_response
             retval = self.fb.make_request(*ARGS, **KWARGS)
@@ -97,6 +99,7 @@ def test_make_request_202(self):
         mock_response = mock.Mock()
         mock_response.status_code = 202
         mock_response.content = "1"
+        mock_response.headers = {}
         ARGS = (1, 2)
         KWARGS = {'a': 3, 'b': 4, 'Accept-Language': self.fb.system}
         with mock.patch.object(self.fb.client, 'make_request') as client_make_request:
@@ -110,6 +113,7 @@ def test_make_request_delete_204(self):
         mock_response = mock.Mock()
         mock_response.status_code = 204
         mock_response.content = "1"
+        mock_response.headers = {}
         ARGS = (1, 2)
         KWARGS = {'a': 3, 'b': 4, 'method': 'DELETE', 'Accept-Language': self.fb.system}
         with mock.patch.object(self.fb.client, 'make_request') as client_make_request:
@@ -123,13 +127,69 @@ def test_make_request_delete_not_204(self):
         mock_response = mock.Mock()
         mock_response.status_code = 205
         mock_response.content = "1"
+        mock_response.headers = {}
         ARGS = (1, 2)
         KWARGS = {'a': 3, 'b': 4, 'method': 'DELETE', 'Accept-Language': self.fb.system}
         with mock.patch.object(self.fb.client, 'make_request') as client_make_request:
             client_make_request.return_value = mock_response
             self.assertRaises(DeleteError, self.fb.make_request, *ARGS, **KWARGS)
 
 
+class RateLimitTest(TestBase):
+    """
+    Test how make_request interacts with Fitbit API's rate-limiting headers
+    """
+
+    def test_updates_parameters_from_request(self):
+        mock_response = mock.Mock()
+        mock_response.status_code = 200
+        mock_response.content = b"1"
+        mock_response.headers = {
+            'fitbit-rate-limit-limit': '150',
+            'fitbit-rate-limit-remaining': '149',
+            'fitbit-rate-limit-reset': '1801',
+        }
+        with mock.patch.object(self.fb.client, 'make_request') as client_make_request:
+            client_make_request.return_value = mock_response
+            self.fb.make_request('x')
+        self.assertEqual(150, self.fb.rate_limit_limit)
+        self.assertEqual(149, self.fb.rate_limit_remaining)
+        self.assertAlmostEqual(time.time() + 1801, self.fb.rate_limit_reset, places=0)
+
+    def test_refuses_requests_that_will_be_throttled(self):
+        mock_response = mock.Mock()
+        mock_response.status_code = 200
+        mock_response.content = b"1"
+        mock_response.headers = {}
+
+        self.fb.rate_limit_limit = 150
+        self.fb.rate_limit_reset = time.time() + 100
+        self.fb.rate_limit_remaining = 1
+
+        # Happy path where we shouldn't be rejected
+        with mock.patch.object(self.fb.client, 'make_request') as client_make_request:
+            client_make_request.return_value = mock_response
+            self.fb.make_request('x')
+
+        # Failure case where we expect a rejection
+        self.fb.rate_limit_remaining = 0
+        try:
+            self.fb.make_request('x')
+            self.fail("Expected fitbit.exceptions.RateLimited to be thrown")
+        except RateLimited as rl:
+            self.assertEqual(150, rl.rate_limit_limit)
+            self.assertEqual(0, rl.rate_limit_remaining)
+            self.assertAlmostEqual(time.time() + 100, rl.rate_limit_reset, places=0)
+
+        # Happy path where remaining is zero, but reset is in the past
+        self.fb.rate_limit_reset = time.time() - 100
+        self.fb.rate_limit_remaining = 1
+
+        with mock.patch.object(self.fb.client, 'make_request') as client_make_request:
+            client_make_request.return_value = mock_response
+            self.fb.make_request('x')
+
+
 class CollectionResourceTest(TestBase):
     """ Tests for _COLLECTION_RESOURCE """
     def test_all_args(self):

fitbit_tests/test_exceptions.py

@@ -27,6 +27,7 @@ def test_response_ok(self):
         r = mock.Mock(spec=requests.Response)
         r.status_code = 200
         r.content = b'{"normal": "resource"}'
+        r.headers = {}
 
         f = Fitbit(**self.client_kwargs)
         f.client._request = lambda *args, **kwargs: r
@@ -73,6 +74,7 @@ def test_response_error(self):
         """
         r = mock.Mock(spec=requests.Response)
         r.content = b'{"normal": "resource"}'
+        r.headers = {}
 
         self.client_kwargs['oauth2'] = True
         f = Fitbit(**self.client_kwargs)
@@ -116,6 +118,7 @@ def test_serialization(self):
         r = mock.Mock(spec=requests.Response)
         r.status_code = 200
         r.content = b"iyam not jason"
+        r.headers = {}
 
         f = Fitbit(**self.client_kwargs)
         f.client._request = lambda *args, **kwargs: r
@@ -128,6 +131,7 @@ def test_delete_error(self):
         r = mock.Mock(spec=requests.Response)
         r.status_code = 201
         r.content = b'{"it\'s all": "ok"}'
+        r.headers = {}
 
         f = Fitbit(**self.client_kwargs)
         f.client._request = lambda *args, **kwargs: r