Wrap UnleashClient directly

Author: cmanallen

sentry_sdk/integrations/unleash.py

@@ -1,43 +1,30 @@
 from functools import wraps
-from typing import TYPE_CHECKING
+from typing import Any
 
 import sentry_sdk
 from sentry_sdk.flag_utils import flag_error_processor
 from sentry_sdk.integrations import Integration, DidNotEnable
 
-if TYPE_CHECKING:
-    from typing import Any, Optional
-
-    try:
-        from UnleashClient import UnleashClient
-    except ImportError:
-        raise DidNotEnable("UnleashClient is not installed")
+try:
+    from UnleashClient import UnleashClient
+except ImportError:
+    raise DidNotEnable("UnleashClient is not installed")
 
 
 class UnleashIntegration(Integration):
     identifier = "unleash"
-    _unleash_client = None  # type: Optional[UnleashClient]
-
-    def __init__(self, unleash_client):
-        # type: (UnleashClient) -> None
-        self.__class__._unleash_client = unleash_client
 
     @staticmethod
     def setup_once():
         # type: () -> None
-
-        client = UnleashIntegration._unleash_client
-        if not client:
-            raise DidNotEnable("Error getting UnleashClient instance")
-
         # Wrap and patch evaluation methods (instance methods)
-        old_is_enabled = client.is_enabled
-        old_get_variant = client.get_variant
+        old_is_enabled = UnleashClient.is_enabled
+        old_get_variant = UnleashClient.get_variant
 
         @wraps(old_is_enabled)
-        def sentry_is_enabled(feature, *a, **kw):
-            # type: (str, *Any, **Any) -> Any
-            enabled = old_is_enabled(feature, *a, **kw)
+        def sentry_is_enabled(self, feature, *args, **kwargs):
+            # type: (UnleashClient, str, *Any, **Any) -> Any
+            enabled = old_is_enabled(self, feature, *args, **kwargs)
 
             # We have no way of knowing what type of unleash feature this is, so we have to treat
             # it as a boolean / toggle feature.
@@ -47,11 +34,10 @@ def sentry_is_enabled(feature, *a, **kw):
             return enabled
 
         @wraps(old_get_variant)
-        def sentry_get_variant(feature, *a, **kw):
-            # type: (str, *Any, **Any) -> Any
-            variant = old_get_variant(feature, *a, **kw)
+        def sentry_get_variant(self, feature, *args, **kwargs):
+            # type: (UnleashClient, str, *Any, **Any) -> Any
+            variant = old_get_variant(self, feature, *args, **kwargs)
             enabled = variant.get("enabled", False)
-            # _payload_type = variant.get("payload", {}).get("type")
 
             # Payloads are not always used as the feature's value for application logic. They
             # may be used for metrics or debugging context instead. Therefore, we treat every
@@ -60,8 +46,8 @@ def sentry_get_variant(feature, *a, **kw):
             flags.set(feature, enabled)
             return variant
 
-        client.is_enabled = sentry_is_enabled  # type: ignore
-        client.get_variant = sentry_get_variant  # type: ignore
+        UnleashClient.is_enabled = sentry_is_enabled  # type: ignore
+        UnleashClient.get_variant = sentry_get_variant  # type: ignore
 
         # Error processor
         scope = sentry_sdk.get_current_scope()

tests/integrations/unleash/test_unleash.py

@@ -2,22 +2,24 @@
 import sys
 from random import random
 from unittest import mock
+from UnleashClient import UnleashClient
 
 import pytest
 
 import sentry_sdk
 from sentry_sdk.integrations.unleash import UnleashIntegration
-from tests.integrations.unleash.testutils import MockUnleashClient
+from tests.integrations.unleash.testutils import mock_unleash_client
 
 
 def test_is_enabled(sentry_init, capture_events, uninstall_integration):
-    client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
 
-    client.is_enabled("hello")
-    client.is_enabled("world")
-    client.is_enabled("other")
+    with mock_unleash_client():
+        client = UnleashClient()
+        sentry_init(integrations=[UnleashIntegration()])
+        client.is_enabled("hello")
+        client.is_enabled("world")
+        client.is_enabled("other")
 
     events = capture_events()
     sentry_sdk.capture_exception(Exception("something wrong!"))
@@ -33,16 +35,17 @@ def test_is_enabled(sentry_init, capture_events, uninstall_integration):
 
 
 def test_get_variant(sentry_init, capture_events, uninstall_integration):
-    client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
 
-    client.get_variant("no_payload_feature")
-    client.get_variant("string_feature")
-    client.get_variant("json_feature")
-    client.get_variant("csv_feature")
-    client.get_variant("number_feature")
-    client.get_variant("unknown_feature")
+    with mock_unleash_client():
+        client = UnleashClient()
+        sentry_init(integrations=[UnleashIntegration()])  # type: ignore
+        client.get_variant("no_payload_feature")
+        client.get_variant("string_feature")
+        client.get_variant("json_feature")
+        client.get_variant("csv_feature")
+        client.get_variant("number_feature")
+        client.get_variant("unknown_feature")
 
     events = capture_events()
     sentry_sdk.capture_exception(Exception("something wrong!"))
@@ -61,25 +64,27 @@ def test_get_variant(sentry_init, capture_events, uninstall_integration):
 
 
 def test_is_enabled_threaded(sentry_init, capture_events, uninstall_integration):
-    client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
-    events = capture_events()
 
-    def task(flag_key):
-        # Creates a new isolation scope for the thread.
-        # This means the evaluations in each task are captured separately.
-        with sentry_sdk.isolation_scope():
-            client.is_enabled(flag_key)
-            # use a tag to identify to identify events later on
-            sentry_sdk.set_tag("task_id", flag_key)
-            sentry_sdk.capture_exception(Exception("something wrong!"))
+    with mock_unleash_client():
+        client = UnleashClient()
+        sentry_init(integrations=[UnleashIntegration()])  # type: ignore
+        events = capture_events()
+
+        def task(flag_key):
+            # Creates a new isolation scope for the thread.
+            # This means the evaluations in each task are captured separately.
+            with sentry_sdk.isolation_scope():
+                client.is_enabled(flag_key)
+                # use a tag to identify to identify events later on
+                sentry_sdk.set_tag("task_id", flag_key)
+                sentry_sdk.capture_exception(Exception("something wrong!"))
 
-    # Capture an eval before we split isolation scopes.
-    client.is_enabled("hello")
+        # Capture an eval before we split isolation scopes.
+        client.is_enabled("hello")
 
-    with cf.ThreadPoolExecutor(max_workers=2) as pool:
-        pool.map(task, ["world", "other"])
+        with cf.ThreadPoolExecutor(max_workers=2) as pool:
+            pool.map(task, ["world", "other"])
 
     # Capture error in original scope
     sentry_sdk.set_tag("task_id", "0")
@@ -108,25 +113,27 @@ def task(flag_key):
 
 
 def test_get_variant_threaded(sentry_init, capture_events, uninstall_integration):
-    client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
-    events = capture_events()
 
-    def task(flag_key):
-        # Creates a new isolation scope for the thread.
-        # This means the evaluations in each task are captured separately.
-        with sentry_sdk.isolation_scope():
-            client.get_variant(flag_key)
-            # use a tag to identify to identify events later on
-            sentry_sdk.set_tag("task_id", flag_key)
-            sentry_sdk.capture_exception(Exception("something wrong!"))
+    with mock_unleash_client():
+        client = UnleashClient()
+        sentry_init(integrations=[UnleashIntegration()])  # type: ignore
+        events = capture_events()
 
-    # Capture an eval before we split isolation scopes.
-    client.get_variant("hello")
+        def task(flag_key):
+            # Creates a new isolation scope for the thread.
+            # This means the evaluations in each task are captured separately.
+            with sentry_sdk.isolation_scope():
+                client.get_variant(flag_key)
+                # use a tag to identify to identify events later on
+                sentry_sdk.set_tag("task_id", flag_key)
+                sentry_sdk.capture_exception(Exception("something wrong!"))
 
-    with cf.ThreadPoolExecutor(max_workers=2) as pool:
-        pool.map(task, ["no_payload_feature", "other"])
+        # Capture an eval before we split isolation scopes.
+        client.get_variant("hello")
+
+        with cf.ThreadPoolExecutor(max_workers=2) as pool:
+            pool.map(task, ["no_payload_feature", "other"])
 
     # Capture error in original scope
     sentry_sdk.set_tag("task_id", "0")
@@ -157,26 +164,27 @@ def task(flag_key):
 @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python3.7 or higher")
 def test_is_enabled_asyncio(sentry_init, capture_events, uninstall_integration):
     asyncio = pytest.importorskip("asyncio")
-
-    client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
-    events = capture_events()
 
-    async def task(flag_key):
-        with sentry_sdk.isolation_scope():
-            client.is_enabled(flag_key)
-            # use a tag to identify to identify events later on
-            sentry_sdk.set_tag("task_id", flag_key)
-            sentry_sdk.capture_exception(Exception("something wrong!"))
+    with mock_unleash_client():
+        client = UnleashClient()
+        sentry_init(integrations=[UnleashIntegration()])  # type: ignore
+        events = capture_events()
 
-    async def runner():
-        return asyncio.gather(task("world"), task("other"))
+        async def task(flag_key):
+            with sentry_sdk.isolation_scope():
+                client.is_enabled(flag_key)
+                # use a tag to identify to identify events later on
+                sentry_sdk.set_tag("task_id", flag_key)
+                sentry_sdk.capture_exception(Exception("something wrong!"))
 
-    # Capture an eval before we split isolation scopes.
-    client.is_enabled("hello")
+        async def runner():
+            return asyncio.gather(task("world"), task("other"))
 
-    asyncio.run(runner())
+        # Capture an eval before we split isolation scopes.
+        client.is_enabled("hello")
+
+        asyncio.run(runner())
 
     # Capture error in original scope
     sentry_sdk.set_tag("task_id", "0")
@@ -208,25 +216,27 @@ async def runner():
 def test_get_variant_asyncio(sentry_init, capture_events, uninstall_integration):
     asyncio = pytest.importorskip("asyncio")
 
-    client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
-    events = capture_events()
 
-    async def task(flag_key):
-        with sentry_sdk.isolation_scope():
-            client.get_variant(flag_key)
-            # use a tag to identify to identify events later on
-            sentry_sdk.set_tag("task_id", flag_key)
-            sentry_sdk.capture_exception(Exception("something wrong!"))
+    with mock_unleash_client():
+        client = UnleashClient()
+        sentry_init(integrations=[UnleashIntegration()])  # type: ignore
+        events = capture_events()
+
+        async def task(flag_key):
+            with sentry_sdk.isolation_scope():
+                client.get_variant(flag_key)
+                # use a tag to identify to identify events later on
+                sentry_sdk.set_tag("task_id", flag_key)
+                sentry_sdk.capture_exception(Exception("something wrong!"))
 
-    async def runner():
-        return asyncio.gather(task("no_payload_feature"), task("other"))
+        async def runner():
+            return asyncio.gather(task("no_payload_feature"), task("other"))
 
-    # Capture an eval before we split isolation scopes.
-    client.get_variant("hello")
+        # Capture an eval before we split isolation scopes.
+        client.get_variant("hello")
 
-    asyncio.run(runner())
+        asyncio.run(runner())
 
     # Capture error in original scope
     sentry_sdk.set_tag("task_id", "0")
@@ -254,39 +264,17 @@ async def runner():
     }
 
 
-def test_client_isolation(sentry_init, capture_events, uninstall_integration):
-    """
-    If the integration is tracking a single client, evaluations from other clients should not be
-    captured.
-    """
-    client = MockUnleashClient()
-    uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
-
-    other_client = MockUnleashClient()
-
-    other_client.is_enabled("hello")
-    other_client.is_enabled("world")
-    other_client.is_enabled("other")
-    other_client.get_variant("no_payload_feature")
-    other_client.get_variant("json_feature")
-
-    events = capture_events()
-    sentry_sdk.capture_exception(Exception("something wrong!"))
-
-    assert len(events) == 1
-    assert events[0]["contexts"]["flags"] == {"values": []}
-
-
 def test_wraps_original(sentry_init, uninstall_integration):
-    client = MockUnleashClient()
-    mock_is_enabled = mock.Mock(return_value=random() < 0.5)
-    client.is_enabled = mock_is_enabled
-    mock_get_variant = mock.Mock(return_value={"enabled": random() < 0.5})
-    client.get_variant = mock_get_variant
+    with mock_unleash_client():
+        client = UnleashClient()
 
-    uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
+        mock_is_enabled = mock.Mock(return_value=random() < 0.5)
+        mock_get_variant = mock.Mock(return_value={"enabled": random() < 0.5})
+        client.is_enabled = mock_is_enabled
+        client.get_variant = mock_get_variant
+
+        uninstall_integration(UnleashIntegration)
+        sentry_init(integrations=[UnleashIntegration()])  # type: ignore
 
     res = client.is_enabled("test-flag", "arg", kwarg=1)
     assert res == mock_is_enabled.return_value
@@ -304,15 +292,17 @@ def test_wraps_original(sentry_init, uninstall_integration):
 
 
 def test_wrapper_attributes(sentry_init, uninstall_integration):
-    client = MockUnleashClient()
-    original_is_enabled = client.is_enabled
-    original_get_variant = client.get_variant
+    with mock_unleash_client():
+        client = UnleashClient()  # <- Returns a MockUnleashClient
 
-    uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
+        original_is_enabled = client.is_enabled
+        original_get_variant = client.get_variant
 
-    assert client.is_enabled.__name__ == "is_enabled"
-    assert client.is_enabled.__qualname__ == original_is_enabled.__qualname__
+        uninstall_integration(UnleashIntegration)
+        sentry_init(integrations=[UnleashIntegration()])  # type: ignore
 
-    assert client.get_variant.__name__ == "get_variant"
-    assert client.get_variant.__qualname__ == original_get_variant.__qualname__
+        # Mock clients methods have not lost their qualified names after decoration.
+        assert client.is_enabled.__name__ == "is_enabled"
+        assert client.is_enabled.__qualname__ == original_is_enabled.__qualname__
+        assert client.get_variant.__name__ == "get_variant"
+        assert client.get_variant.__qualname__ == original_get_variant.__qualname__