Merge pull request #11 from aropan/fix-too-many-requests

Fix too many requests

pymailinator/tests/test_wrapper.py

@@ -45,6 +45,18 @@ def get_message(url):
     return MockResponse(status=200, filename=filename, url=url)
 
 
+class MockServerTooManyRequests(object):
+    def __init__(self, count, response):
+        self.count = count
+        self.response = response
+
+    def get(self, url):
+        if self.count > 0:
+            self.count -= 1
+            return MockResponse(status=429, content='', url=url)
+        return self.response(url)
+
+
 # noinspection PyArgumentList
 class TestWrapper(unittest.TestCase):
     def test_empty_mailbox(self):
@@ -70,72 +82,99 @@ def test_invalid_token(self):
         with self.assertRaises(wrapper.MissingToken):
             inbox.get()
 
-    def test_successful_mailbox(self):
-        wrapper.urlopen = get_mailbox
+    def test_successful_mailbox(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
         self.assertGreater(inbox.count(), 0)
 
-    def test_successful_message(self):
-        wrapper.urlopen = get_mailbox
+    def test_successful_message(self, mailbox=get_mailbox, message=get_message):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
-        wrapper.urlopen = get_message
+        wrapper.urlopen = message
         message = inbox.messages[0]
         message.get_message()
         self.assertNotEqual(message.body, '')
 
-    def test_missing_message(self):
-        wrapper.urlopen = get_mailbox
+    def test_missing_message(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
         wrapper.urlopen = get_missing_message
         with self.assertRaises(wrapper.MessageNotFound):
             inbox.messages[0].get_message()
 
-    def test_view_subjects(self):
-        wrapper.urlopen = get_mailbox
+    def test_view_subjects(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
         self.assertEqual(type(inbox.view_subjects()), list)
         self.assertGreater(len(inbox.view_subjects()), 0)
 
-    def test_view_message_ids(self):
-        wrapper.urlopen = get_mailbox
+    def test_view_message_ids(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
         self.assertEqual(type(inbox.view_message_ids()), list)
         self.assertGreater(len(inbox.view_message_ids()), 0)
 
-    def test_get_message_by_subject(self):
-        wrapper.urlopen = get_mailbox
+    def test_get_message_by_subject(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
         message = inbox.get_message_by_subject('Want to cheat? ')
         self.assertEqual(message.subject, 'Want to cheat? ')
 
-    def test_get_message_by_id(self):
-        wrapper.urlopen = get_mailbox
+    def test_get_message_by_id(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
         message = inbox.get_message_by_id('1418740612-3134545-m8r-rmtci4')
         self.assertEqual(message.id, '1418740612-3134545-m8r-rmtci4')
 
-    def test_filter_inbox(self):
-        wrapper.urlopen = get_mailbox
+    def test_filter_inbox(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get()
         filtered_me = inbox.filter('to', 'me')
         self.assertEqual(len(filtered_me), 0)
         filtered = inbox.filter('to', 'm8r-rmtci4@mailinator.com')
         self.assertEqual(len(filtered), 2)
 
-    def test_other_mailbox(self):
-        wrapper.urlopen = get_mailbox
+    def test_other_mailbox(self, mailbox=get_mailbox):
+        wrapper.urlopen = mailbox
         inbox = wrapper.Inbox('123')
         inbox.get('other')
         self.assertGreater(inbox.count(), 0)
 
+    def tests_too_many_request(self):
+        wrapper.TIME_SLEEP_TOO_MANY_REQUESTS = .1
+        for test_func in (
+            'test_successful_mailbox',
+            'test_successful_message',
+            'test_missing_message',
+            'test_view_subjects',
+            'test_view_message_ids',
+            'test_get_message_by_subject',
+            'test_get_message_by_id',
+            'test_filter_inbox',
+            'test_other_mailbox',
+        ):
+            mock_server = MockServerTooManyRequests(2, get_mailbox)
+            getattr(self, test_func)(mailbox=mock_server.get)
+
+        for test_func in (
+            'test_successful_message',
+        ):
+            mock_server = MockServerTooManyRequests(2, get_message)
+            getattr(self, test_func)(message=mock_server.get)
+
+        with self.assertRaises(wrapper.TooManyRequests):
+            wrapper.NUMBER_ATTEMPTS_TOO_MANY_REQUESTS = 3
+            mock_server = MockServerTooManyRequests(4, get_mailbox)
+            self.test_successful_mailbox(mailbox=mock_server.get)
+
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()

pymailinator/wrapper.py

@@ -1,4 +1,5 @@
 import json
+from time import sleep
 
 try:
     from urllib.request import urlopen
@@ -7,6 +8,14 @@
     from urllib import urlopen, urlencode
 
 
+TIME_SLEEP_TOO_MANY_REQUESTS = .7
+NUMBER_ATTEMPTS_TOO_MANY_REQUESTS = 12
+
+
+class TooManyRequests(Exception):
+    pass
+
+
 class InvalidToken(Exception):
     pass
 
@@ -45,7 +54,7 @@ def __init__(self, token, data):
     def get_message(self):
         query_string = {'token': self.token, 'msgid': self.id}
         url = self._baseURL + "?" + urlencode(query_string)
-        request = urlopen(url)
+        request = get_request(url)
         if request.getcode() == 404:
             raise MessageNotFound
         response = request.read()
@@ -89,7 +98,7 @@ def get(self, mailbox=None, private_domain=False):
         if private_domain:
             query_string.update({'private_domain': json.dumps(private_domain)})
         url = self._baseURL + '?' + urlencode(query_string)
-        request = urlopen(url)
+        request = get_request(url)
         if request.getcode() == 400:
             raise InvalidToken
         response = request.read()
@@ -144,3 +153,13 @@ def clean_response(response):
     else:
         return response
 
+
+def get_request(url, sleep_time=.7):
+    for i in range(NUMBER_ATTEMPTS_TOO_MANY_REQUESTS):
+        request = urlopen(url)
+        if request.getcode() != 429:
+            break
+        sleep(TIME_SLEEP_TOO_MANY_REQUESTS)
+    else:
+        raise TooManyRequests
+    return request