Refactor response validation and improve rate limit detection

onegram/exceptions.py

@@ -14,6 +14,12 @@ class AuthUserError(AuthException):
 class NotSupportedError(OnegramException):
     pass
 
+class RequestFailed(OnegramException):
+    pass
+
+class RateLimitedError(RequestFailed):
+    pass
+
 
 # TODO [romeira]: Query/action exceptions {06/03/18 23:08}
 # TODO [romeira]: Session expired exception {06/03/18 23:08}

onegram/session.py

@@ -19,8 +19,9 @@
 from .constants import URLS
 from .constants import REGEXES
 from .exceptions import AuthException, NotSupportedError
+from .exceptions import RateLimitedError
 from .utils.ratelimit import RateLimiter
-from .utils.validation import check_auth
+from .utils.validation import validate_response
 
 
 class _BaseSession(Session):
@@ -103,20 +104,13 @@ def _after_request_attempt(func, trial_number, *a, **kw):
             self.logger.warning(f'RETRY {trial_number} attempt(s) ...')
 
         @retry(wait=wait_chain(wait_fixed(60), wait_fixed(15)),
-               retry=retry_if_exception_type(HTTPError),
+               retry=retry_if_exception_type(RateLimitedError),
                after=_after_request_attempt)
         def _request():
             with self.rate_limiter:
                 self.logger.info(f'{method} "{url}"')
-                response = None
-                try:
-                    response = self._requests.request(method, url, *a, **kw)
-                    response.raise_for_status()
-                    return json.loads(response.text)
-                except Exception:
-                    if response:
-                        self.logger.error(response.text)
-                    raise
+                response = self._requests.request(method, url, *a, **kw)
+                return validate_response(self, response)
         return _request()
 
 
@@ -181,11 +175,9 @@ def _login(self):
         headers = ACTION_HEADERS
         headers['X-CSRFToken'] = self.cookies['csrftoken']
         kw['headers'] = headers
-
         response = self._requests.post(URLS['login'], **kw)
-        response.raise_for_status()
-        check_auth(json.loads(response.text))
-        self.user_id = self.cookies.get('ds_user_id')
+        self.user_id = self.cookies.get('ds_user_id', None)
+        return validate_response(self, response, auth=True)
 
 
     def __str__(self):

onegram/utils/validation.py

@@ -1,8 +1,46 @@
+import json
 from ..exceptions import AuthFailed, AuthUserError
+from ..exceptions import RequestFailed, RateLimitedError
 
 
-def check_auth(response):
+def validate_response(session, response, auth=False):
+    try:
+        try:
+            js_response = json.loads(response.text)
+        except:
+            response.raise_for_status()
+            if auth:
+                raise AuthFailed('Authentication failed')
+        else:
+            if auth:
+                _check_auth(js_response)
+            _check_status(js_response)
+            response.raise_for_status()
+            return js_response
+    except:
+        if response:
+            session.logger.error(response.text)
+        raise
+
+
+def _check_auth(response):
     if not response.get('user', False):
         raise AuthUserError('Please check your username')
     if not response.get('authenticated', False):
         raise AuthFailed('Authentication failed')
+
+
+def _check_status(response):
+    message = response.get('message', '').lower()
+    status = response.get('status', '').lower()
+
+    msg = ''
+    if status:
+        msg = status
+    if message:
+        msg += f': {message}'
+
+    if 'rate limit' in message:
+        raise RateLimitedError(msg)
+    if 'fail' in status:
+        raise RequestFailed(msg)