fix: Test transport rate limits parsing and enforcement (getsentry#652)

untitaker · web-flow · commit b7679a50f31e · 2020-03-23T17:24:22.000+01:00
Also fix a bug where missing categories ("123::project") would not
enforce a rate limit for all categories, as they were parsed as category
"" instead of category None.
diff --git a/sentry_sdk/transport.py b/sentry_sdk/transport.py
@@ -15,17 +15,22 @@
 from sentry_sdk._types import MYPY
 
 if MYPY:
-    from typing import Type
     from typing import Any
-    from typing import Optional
+    from typing import Callable
     from typing import Dict
+    from typing import Iterable
+    from typing import Optional
+    from typing import Tuple
+    from typing import Type
     from typing import Union
-    from typing import Callable
+
     from urllib3.poolmanager import PoolManager  # type: ignore
     from urllib3.poolmanager import ProxyManager
 
     from sentry_sdk._types import Event
 
+    DataCategory = Optional[str]
+
 try:
     from urllib.request import getproxies
 except ImportError:
@@ -94,6 +99,21 @@ def __del__(self):
             pass
 
 
+def _parse_rate_limits(header, now=None):
+    # type: (Any, Optional[datetime]) -> Iterable[Tuple[DataCategory, datetime]]
+    if now is None:
+        now = datetime.utcnow()
+
+    for limit in header.split(","):
+        try:
+            retry_after, categories, _ = limit.strip().split(":", 2)
+            retry_after = now + timedelta(seconds=int(retry_after))
+            for category in categories and categories.split(";") or (None,):
+                yield category, retry_after
+        except (LookupError, ValueError):
+            continue
+
+
 class HttpTransport(Transport):
     """The default HTTP transport."""
 
@@ -107,7 +127,7 @@ def __init__(
         assert self.parsed_dsn is not None
         self._worker = BackgroundWorker()
         self._auth = self.parsed_dsn.to_auth("sentry.python/%s" % VERSION)
-        self._disabled_until = {}  # type: Dict[Any, datetime]
+        self._disabled_until = {}  # type: Dict[DataCategory, datetime]
         self._retry = urllib3.util.Retry()
         self.options = options
 
@@ -129,16 +149,7 @@ def _update_rate_limits(self, response):
         # no matter of the status code to update our internal rate limits.
         header = response.headers.get("x-sentry-rate-limit")
         if header:
-            for limit in header.split(","):
-                try:
-                    retry_after, categories, _ = limit.strip().split(":", 2)
-                    retry_after = datetime.utcnow() + timedelta(
-                        seconds=int(retry_after)
-                    )
-                    for category in categories.split(";") or (None,):
-                        self._disabled_until[category] = retry_after
-                except (LookupError, ValueError):
-                    continue
+            self._disabled_until.update(_parse_rate_limits(header))
 
         # old sentries only communicate global rate limit hits via the
         # retry-after header on 429.  This header can also be emitted on new
diff --git a/tests/test_transport.py b/tests/test_transport.py
@@ -2,11 +2,12 @@
 import logging
 import pickle
 
-from datetime import datetime
+from datetime import datetime, timedelta
 
 import pytest
 
 from sentry_sdk import Hub, Client, add_breadcrumb, capture_message
+from sentry_sdk.transport import _parse_rate_limits
 
 
 @pytest.fixture(params=[True, False])
@@ -54,3 +55,115 @@ def test_transport_works(
     assert httpserver.requests
 
     assert any("Sending event" in record.msg for record in caplog.records) == debug
+
+
+NOW = datetime(2014, 6, 2)
+
+
+@pytest.mark.parametrize(
+    "input,expected",
+    [
+        # Invalid rate limits
+        ("", {}),
+        ("invalid", {}),
+        (",,,", {}),
+        (
+            "42::organization, invalid, 4711:foobar;transaction;security:project",
+            {
+                None: NOW + timedelta(seconds=42),
+                "transaction": NOW + timedelta(seconds=4711),
+                "security": NOW + timedelta(seconds=4711),
+                # Unknown data categories
+                "foobar": NOW + timedelta(seconds=4711),
+            },
+        ),
+        (
+            "4711:foobar;;transaction:organization",
+            {
+                "transaction": NOW + timedelta(seconds=4711),
+                # Unknown data categories
+                "foobar": NOW + timedelta(seconds=4711),
+                "": NOW + timedelta(seconds=4711),
+            },
+        ),
+    ],
+)
+def test_parse_rate_limits(input, expected):
+    assert dict(_parse_rate_limits(input, now=NOW)) == expected
+
+
+def test_simple_rate_limits(httpserver, capsys, caplog):
+    client = Client(dsn="http://foobar@{}/123".format(httpserver.url[len("http://") :]))
+    httpserver.serve_content("no", 429, headers={"Retry-After": "4"})
+
+    client.capture_event({"type": "transaction"})
+    client.flush()
+
+    assert len(httpserver.requests) == 1
+    del httpserver.requests[:]
+
+    assert set(client.transport._disabled_until) == set([None])
+
+    client.capture_event({"type": "transaction"})
+    client.capture_event({"type": "event"})
+    client.flush()
+
+    assert not httpserver.requests
+
+
+@pytest.mark.parametrize("response_code", [200, 429])
+def test_data_category_limits(httpserver, capsys, caplog, response_code):
+    client = Client(
+        dict(dsn="http://foobar@{}/123".format(httpserver.url[len("http://") :]))
+    )
+    httpserver.serve_content(
+        "hm",
+        response_code,
+        headers={"X-Sentry-Rate-Limit": "4711:transaction:organization"},
+    )
+
+    client.capture_event({"type": "transaction"})
+    client.flush()
+
+    assert len(httpserver.requests) == 1
+    del httpserver.requests[:]
+
+    assert set(client.transport._disabled_until) == set(["transaction"])
+
+    client.transport.capture_event({"type": "transaction"})
+    client.transport.capture_event({"type": "transaction"})
+    client.flush()
+
+    assert not httpserver.requests
+
+    client.capture_event({"type": "event"})
+    client.flush()
+
+    assert len(httpserver.requests) == 1
+
+
+@pytest.mark.parametrize("response_code", [200, 429])
+def test_complex_limits_without_data_category(
+    httpserver, capsys, caplog, response_code
+):
+    client = Client(
+        dict(dsn="http://foobar@{}/123".format(httpserver.url[len("http://") :]))
+    )
+    httpserver.serve_content(
+        "hm", response_code, headers={"X-Sentry-Rate-Limit": "4711::organization"},
+    )
+
+    client.capture_event({"type": "transaction"})
+    client.flush()
+
+    assert len(httpserver.requests) == 1
+    del httpserver.requests[:]
+
+    assert set(client.transport._disabled_until) == set([None])
+
+    client.transport.capture_event({"type": "transaction"})
+    client.transport.capture_event({"type": "transaction"})
+    client.capture_event({"type": "event"})
+    client.flush()
+
+    assert len(httpserver.requests) == 0
