diff --git a/tests/test_rate_limit.py b/tests/test_rate_limit.py index 6edd7e0c..5b2c32fe 100644 --- a/tests/test_rate_limit.py +++ b/tests/test_rate_limit.py @@ -1,5 +1,6 @@ # encoding: utf-8 +import time import re import sys import unittest @@ -9,7 +10,11 @@ import responses warnings.filterwarnings('ignore', category=DeprecationWarning) -DEF_URL_RE = re.compile(r'https?://.*\.twitter.com/1\.1/.*') +DEFAULT_URL = re.compile(r'https?://.*\.twitter.com/1\.1/.*') + +HEADERS = {'x-rate-limit-limit': '63', + 'x-rate-limit-remaining': '63', + 'x-rate-limit-reset': '626672700'} class ErrNull(object): @@ -60,7 +65,7 @@ def testInitializeRateLimit(self): self.assertTrue(self.api.rate_limit) self.assertTrue(self.api.sleep_on_rate_limit) - responses.add(responses.GET, url=DEF_URL_RE, body=b'{}', status=200) + responses.add(responses.GET, url=DEFAULT_URL, body=b'{}', status=200) try: self.api.GetStatus(status_id=1234) self.api.GetUser(screen_name='test') @@ -112,7 +117,6 @@ def setUp(self): self.api.InitializeRateLimit() self.assertTrue(self.api.rate_limit) - def tearDown(self): sys.stderr = self._stderr pass @@ -124,12 +128,15 @@ def testGetRateLimit(self): self.assertEqual(lim.reset, 1452254278) def testNonStandardEndpointRateLimit(self): - lim = self.api.rate_limit.get_limit('https://api.twitter.com/1.1/geo/id/312.json?skip_status=True') + lim = self.api.rate_limit.get_limit( + 'https://api.twitter.com/1.1/geo/id/312.json?skip_status=True') self.assertEqual(lim.limit, 47) - lim = self.api.rate_limit.get_limit('https://api.twitter.com/1.1/saved_searches/destroy/312.json') + lim = self.api.rate_limit.get_limit( + 'https://api.twitter.com/1.1/saved_searches/destroy/312.json') self.assertEqual(lim.limit, 15) - lim = self.api.rate_limit.get_limit('https://api.twitter.com/1.1/statuses/retweets/312.json?skip_status=True') + lim = self.api.rate_limit.get_limit( + 'https://api.twitter.com/1.1/statuses/retweets/312.json?skip_status=True') self.assertEqual(lim.limit, 23) def testSetRateLimit(self): @@ -157,3 +164,67 @@ def testSetUnknownRateLimit(self): limit = self.api.rate_limit.get_limit( url='https://api.twitter.com/1.1/not/a/real/endpoint.json') self.assertEqual(limit.remaining, 14) + + @responses.activate + def testLimitsViaHeadersNoSleep(self): + api = twitter.Api( + consumer_key='test', + consumer_secret='test', + access_token_key='test', + access_token_secret='test', + sleep_on_rate_limit=False) + + # Get initial rate limit data to populate api.rate_limit object + with open('testdata/ratelimit.json') as f: + resp_data = f.read() + url = '%s/application/rate_limit_status.json' % self.api.base_url + responses.add(responses.GET, url, body=resp_data, match_querystring=True) + + # Add a test URL just to have headers present + responses.add( + method=responses.GET, url=DEFAULT_URL, body='{}', adding_headers=HEADERS) + resp = api.GetSearch(term='test') + self.assertTrue(api.rate_limit.__dict__) + self.assertEqual(api.rate_limit.get_limit('/search/tweets').limit, 63) + self.assertEqual(api.rate_limit.get_limit('/search/tweets').remaining, 63) + self.assertEqual(api.rate_limit.get_limit('/search/tweets').reset, 626672700) + + # No other resource families should be set during above operation. + test_url = '/lists/subscribers/show.json' + self.assertEqual( + api.rate_limit.__dict__.get('resources').get(test_url), None + ) + + # But if we check them, it should go to default 15/15 + self.assertEqual(api.rate_limit.get_limit(test_url).remaining, 15) + self.assertEqual(api.rate_limit.get_limit(test_url).limit, 15) + + @responses.activate + def testLimitsViaHeadersWithSleep(self): + api = twitter.Api( + consumer_key='test', + consumer_secret='test', + access_token_key='test', + access_token_secret='test', + sleep_on_rate_limit=True) + + # Add handler for ratelimit check + url = '%s/application/rate_limit_status.json' % api.base_url + responses.add( + method=responses.GET, url=url, body='{}', match_querystring=True) + + # Get initial rate limit data to populate api.rate_limit object + url = "{0}/search/tweets.json?result_type=mixed&q=test&count=15".format( + api.base_url) + responses.add( + method=responses.GET, + url=url, + body='{}', + match_querystring=True, + adding_headers=HEADERS) + + resp = api.GetSearch(term='test') + self.assertTrue(api.rate_limit) + self.assertEqual(api.rate_limit.get_limit('/search/tweets').limit, 63) + self.assertEqual(api.rate_limit.get_limit('/search/tweets').remaining, 63) + self.assertEqual(api.rate_limit.get_limit('/search/tweets').reset, 626672700) diff --git a/twitter/api.py b/twitter/api.py index 26cf9ddc..17f3fedd 100644 --- a/twitter/api.py +++ b/twitter/api.py @@ -186,7 +186,7 @@ def __init__(self, self._InitializeUserAgent() self._InitializeDefaultParameters() - self.rate_limit = None + self.rate_limit = RateLimit() self.sleep_on_rate_limit = sleep_on_rate_limit if base_url is None: @@ -4673,7 +4673,7 @@ def CheckRateLimit(self, url): namedtuple: EndpointRateLimit namedtuple. """ - if not self.rate_limit: + if not self.rate_limit.__dict__.get('resources', None): self.InitializeRateLimit() if url: @@ -4853,7 +4853,7 @@ def _RequestUrl(self, url, verb, data=None, json=None): else: resp = 0 # if not a POST or GET request - if url and self.sleep_on_rate_limit and self.rate_limit: + if url and self.rate_limit: limit = resp.headers.get('x-rate-limit-limit', 0) remaining = resp.headers.get('x-rate-limit-remaining', 0) reset = resp.headers.get('x-rate-limit-reset', 0) diff --git a/twitter/ratelimit.py b/twitter/ratelimit.py index 45a0c889..2910a13a 100644 --- a/twitter/ratelimit.py +++ b/twitter/ratelimit.py @@ -146,6 +146,12 @@ def set_unknown_limit(self, url, limit, remaining, reset): """ endpoint = self.url_to_resource(url) resource_family = endpoint.split('/')[1] + + # This is a conversative check against a situation in which a user has + # set api.sleep_on_rate_limit to True, but hasn't made a call + if not self.__dict__.get('resources', None): + self.__dict__['resources'] = {} + self.__dict__['resources'].update( {resource_family: { endpoint: {