Patch class method instead of instance, so passing in client is not needed. Ref tests back to 1 file

Author: aliu39

sentry_sdk/integrations/unleash.py

@@ -1,42 +1,31 @@
 from functools import wraps
+from typing import TYPE_CHECKING
 
 import sentry_sdk
 from sentry_sdk.flag_utils import flag_error_processor
-from sentry_sdk.integrations import Integration, DidNotEnable
-
-from typing import TYPE_CHECKING
+from sentry_sdk.integrations import Integration
 
+from UnleashClient import UnleashClient
 
 if TYPE_CHECKING:
-    from typing import Optional, Any
-    from UnleashClient import UnleashClient
+    from typing import Any
 
 
 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
 
         # Wrap and patch evaluation functions
-        client = UnleashIntegration._unleash_client
-        if not client:
-            raise DidNotEnable("Error finding UnleashClient instance.")
-
-        # Instance methods (self is excluded)
-        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, *a, **kw):
+            # type: (UnleashClient, str, *Any, **Any) -> Any
+            enabled = old_is_enabled(self, feature, *a, **kw)
 
             # We have no way of knowing what type of feature this is. Have to treat it as
             # boolean flag. TODO: Unless we fetch a list of non-bool flags on startup..
@@ -46,9 +35,9 @@ 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, *a, **kw):
+            # type: (UnleashClient, str, *Any, **Any) -> Any
+            variant = old_get_variant(self, feature, *a, **kw)
             enabled = variant.get("enabled", False)
             payload_type = variant.get("payload", {}).get("type")
 
@@ -58,8 +47,8 @@ def sentry_get_variant(feature, *a, **kw):
 
             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/__init__.py

@@ -0,0 +1,39 @@
+class MockUnleashClient:
+
+    def __init__(self, *a, **kw):
+        self.features = {
+            "hello": True,
+            "world": False,
+        }
+
+        self.feature_to_variant = {
+            "string_feature": {
+                "name": "variant1",
+                "enabled": True,
+                "payload": {"type": "string", "value": "val1"},
+            },
+            "json_feature": {
+                "name": "variant1",
+                "enabled": True,
+                "payload": {"type": "json", "value": '{"key1": 0.53}'},
+            },
+            "number_feature": {
+                "name": "variant1",
+                "enabled": True,
+                "payload": {"type": "number", "value": "134.5"},
+            },
+            "csv_feature": {
+                "name": "variant1",
+                "enabled": True,
+                "payload": {"type": "csv", "value": "abc 123\ncsbq 94"},
+            },
+            "toggle_feature": {"name": "variant1", "enabled": True},
+        }
+
+        self.disabled_variant = {"name": "disabled", "enabled": False}
+
+    def is_enabled(self, feature, *a, **kw):
+        return self.features.get(feature, False)
+
+    def get_variant(self, feature, *a, **kw):
+        return self.feature_to_variant.get(feature, self.disabled_variant)

tests/integrations/unleash/test_unleash.py

@@ -0,0 +1,54 @@
+from unittest.mock import patch
+
+import sentry_sdk
+from sentry_sdk.integrations.unleash import UnleashIntegration
+from tests.integrations.unleash import MockUnleashClient
+
+
+@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
+def test_is_enabled(sentry_init, capture_events, uninstall_integration):
+    mock_unleash_client = MockUnleashClient()
+
+    uninstall_integration(UnleashIntegration.identifier)
+    sentry_init(integrations=[UnleashIntegration()])
+
+    mock_unleash_client.is_enabled("hello")
+    mock_unleash_client.is_enabled("world")
+    mock_unleash_client.is_enabled("other")
+
+    events = capture_events()
+    sentry_sdk.capture_exception(Exception("something wrong!"))
+
+    assert len(events) == 1
+    assert events[0]["contexts"]["flags"] == {
+        "values": [
+            {"flag": "hello", "result": True},
+            {"flag": "world", "result": False},
+            {"flag": "other", "result": False},
+        ]
+    }
+
+
+@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
+def test_get_variant(sentry_init, capture_events, uninstall_integration):
+    mock_unleash_client = MockUnleashClient()
+    uninstall_integration(UnleashIntegration.identifier)
+    sentry_init(integrations=[UnleashIntegration()])
+
+    mock_unleash_client.get_variant("toggle_feature")
+    mock_unleash_client.get_variant("string_feature")
+    mock_unleash_client.get_variant("json_feature")
+    mock_unleash_client.get_variant("csv_feature")
+    mock_unleash_client.get_variant("number_feature")
+    mock_unleash_client.get_variant("unknown_feature")
+
+    events = capture_events()
+    sentry_sdk.capture_exception(Exception("something wrong!"))
+
+    assert len(events) == 1
+    assert events[0]["contexts"]["flags"] == {
+        "values": [
+            {"flag": "toggle_feature", "result": True},
+            {"flag": "unknown_feature", "result": False},
+        ]
+    }