Add SparkRateLimitWarning and use it for rate-limit test

Add a SparkRateLimitWarning class to warn users when a rate-limit event has been detected, and use this new warning (in place of the logging detection) to detect when a rate-limit event has occured in the rate-limit test.

ciscosparkapi/__init__.py

@@ -27,7 +27,8 @@
 from .api.teams import TeamsAPI as _TeamsAPI
 from .api.webhooks import WebhooksAPI as _WebhooksAPI
 from .exceptions import (
-    SparkApiError, SparkRateLimitError, ciscosparkapiException,
+    SparkApiError, SparkRateLimitError, SparkRateLimitWarning,
+    ciscosparkapiException,
 )
 from .models import (
     AccessToken, License, Membership, Message, Organization, Person, Role,

ciscosparkapi/exceptions.py

@@ -152,3 +152,23 @@ def __init__(self, response):
         present in the response headers.
 
         """
+
+
+class SparkRateLimitWarning(UserWarning):
+    """Cisco Spark rate-limit exceeded warning; the request will be retried."""
+
+    def __init__(self, response):
+        super(SparkRateLimitWarning, self).__init__()
+        self.retry_after = int(response.headers.get('Retry-After', 200))
+        """The `Retry-After` time period (in seconds) provided by Cisco Spark.
+
+        Defaults to 200 seconds if the response `Retry-After` header isn't
+        present in the response headers.
+
+        """
+
+    def __str__(self):
+        """Spark rate-limit exceeded warning message."""
+        return "Rate-limit response received; the request will " \
+               "automatically be retried in {0} seconds." \
+               "".format(self.retry_after)

ciscosparkapi/restsession.py

@@ -21,7 +21,9 @@
 from past.builtins import basestring
 import requests
 
-from .exceptions import SparkRateLimitError, ciscosparkapiException
+from .exceptions import (
+    SparkRateLimitError, SparkRateLimitWarning, ciscosparkapiException
+)
 from .response_codes import EXPECTED_RESPONSE_CODE
 from .utils import (
     check_response_code, extract_and_parse_json, validate_base_url,
@@ -58,24 +60,25 @@ def _fix_next_url(next_url):
 
     Raises:
         AssertionError: If the parameter types are incorrect.
-        ciscosparkapiException: If 'next_url' does not contain a valid API
-            endpoint URL (scheme, netloc and path).
+        ValueError: If 'next_url' does not contain a valid API endpoint URL
+            (scheme, netloc and path).
 
     """
     next_url = str(next_url)
     parsed_url = urllib.parse.urlparse(next_url)
 
     if not parsed_url.scheme or not parsed_url.netloc or not parsed_url.path:
-        error_message = "'next_url' must be a valid API endpoint URL, " \
-                        "minimally containing a scheme, netloc and path."
-        raise ciscosparkapiException(error_message)
+        raise ValueError(
+            "'next_url' must be a valid API endpoint URL, minimally "
+            "containing a scheme, netloc and path."
+        )
 
     if parsed_url.query:
         query_list = parsed_url.query.split('&')
         if 'max=null' in query_list:
             query_list.remove('max=null')
             warnings.warn("`max=null` still present in next-URL returned "
-                          "from Cisco Spark", Warning)
+                          "from Cisco Spark", RuntimeWarning)
         new_query = '&'.join(query_list)
         parsed_url = list(parsed_url)
         parsed_url[4] = new_query
@@ -275,9 +278,7 @@ def request(self, method, url, erc, **kwargs):
 
                 # Wait and retry if automatic rate-limit handling is enabled
                 if self.wait_on_rate_limit and e.retry_after:
-                    logger.info("Received rate-limit message; "
-                                "waiting {0} seconds."
-                                "".format(e.retry_after))
+                    warnings.warn(SparkRateLimitWarning(response))
                     time.sleep(e.retry_after)
                     continue
 

tests/test_restsession.py

@@ -3,31 +3,30 @@
 
 
 import logging
+import warnings
 
 import pytest
 
+import ciscosparkapi
+
 
 __author__ = "Chris Lunsford"
 __author_email__ = "[email protected]"
 __copyright__ = "Copyright (c) 2016-2018 Cisco and/or its affiliates."
 __license__ = "MIT"
 
 
-# Helper Classes
-class RateLimitDetector(logging.Handler):
-    """Detects occurrences of rate limiting."""
-
-    def __init__(self):
-        super(RateLimitDetector, self).__init__()
-
-        self.rate_limit_detected = False
+logging.captureWarnings(True)
 
-    def emit(self, record):
-        """Check record to see if it is a rate-limit message."""
-        assert isinstance(record, logging.LogRecord)
 
-        if "Received rate-limit message" in record.msg:
-            self.rate_limit_detected = True
+# Helper Functions
+def rate_limit_detected(w):
+    """Check to see if a rate-limit warning is in the warnings list."""
+    while w:
+        if issubclass(w.pop().category, ciscosparkapi.SparkRateLimitWarning):
+            return True
+            break
+    return False
 
 
 # CiscoSparkAPI Tests
@@ -36,35 +35,16 @@ class TestRestSession:
 
     @pytest.mark.ratelimit
     def test_rate_limit_retry(self, api, rooms_list, add_rooms):
-        logger = logging.getLogger(__name__)
-
         # Save state and initialize test setup
         original_wait_on_rate_limit = api._session.wait_on_rate_limit
         api._session.wait_on_rate_limit = True
 
-        # Add log handler
-        root_logger = logging.getLogger()
-        rate_limit_detector = RateLimitDetector()
-        root_logger.addHandler(rate_limit_detector)
-        logger = logging.getLogger(__name__)
-        logger.info("Starting Rate Limit Testing")
-
-        try:
-            # Try and trigger a rate-limit
-            rooms = api.rooms.list(max=1)
-            request_count = 0
-            while not rate_limit_detector.rate_limit_detected:
-                for room in rooms:
-                    request_count += 1
-                    if rate_limit_detector.rate_limit_detected:
-                        break
-
-        finally:
-            logger.info("Rate-limit reached with approximately %s requests.",
-                        request_count)
-            # Remove the log handler and restore the pre-test state
-            root_logger.removeHandler(rate_limit_detector)
-            api._session.wait_on_rate_limit = original_wait_on_rate_limit
+        with warnings.catch_warnings(record=True) as w:
+            rooms = api.rooms.list()
+            while True:
+                # Try and trigger a rate-limit
+                list(rooms)
+                if rate_limit_detected(w):
+                    break
 
-        # Assert test condition
-        assert rate_limit_detector.rate_limit_detected == True
+        api._session.wait_on_rate_limit = original_wait_on_rate_limit