Patch an UnleashClient instance instead of the class

Author: aliu39

sentry_sdk/integrations/unleash.py

@@ -5,30 +5,39 @@
 from sentry_sdk.flag_utils import flag_error_processor
 from sentry_sdk.integrations import Integration, DidNotEnable
 
-try:
-    from UnleashClient import UnleashClient
-except ImportError:
-    raise DidNotEnable("UnleashClient is not installed")
-
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Optional
+
+    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: (Optional[UnleashClient]) -> None
+        self.__class__._unleash_client = unleash_client
 
     @staticmethod
     def setup_once():
         # type: () -> None
 
-        # Wrap and patch evaluation functions
-        old_is_enabled = UnleashClient.is_enabled
-        old_get_variant = UnleashClient.get_variant
+        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
 
         @wraps(old_is_enabled)
-        def sentry_is_enabled(self, feature, *a, **kw):
-            # type: (UnleashClient, str, *Any, **Any) -> Any
-            enabled = old_is_enabled(self, feature, *a, **kw)
+        def sentry_is_enabled(feature, *a, **kw):
+            # type: (str, *Any, **Any) -> Any
+            enabled = old_is_enabled(feature, *a, **kw)
 
             # We have no way of knowing what type of unleash feature this is, so we have to treat
             # it as a boolean / toggle feature.
@@ -38,9 +47,9 @@ def sentry_is_enabled(self, feature, *a, **kw):
             return enabled
 
         @wraps(old_get_variant)
-        def sentry_get_variant(self, feature, *a, **kw):
-            # type: (UnleashClient, str, *Any, **Any) -> Any
-            variant = old_get_variant(self, feature, *a, **kw)
+        def sentry_get_variant(feature, *a, **kw):
+            # type: (str, *Any, **Any) -> Any
+            variant = old_get_variant(feature, *a, **kw)
             enabled = variant.get("enabled", False)
             # _payload_type = variant.get("payload", {}).get("type")
 
@@ -51,8 +60,8 @@ def sentry_get_variant(self, feature, *a, **kw):
             flags.set(feature, enabled)
             return variant
 
-        UnleashClient.is_enabled = sentry_is_enabled  # type: ignore
-        UnleashClient.get_variant = sentry_get_variant  # type: ignore
+        client.is_enabled = sentry_is_enabled  # type: ignore
+        client.get_variant = sentry_get_variant  # type: ignore
 
         # Error processor
         scope = sentry_sdk.get_current_scope()

tests/integrations/unleash/test_unleash.py

@@ -1,23 +1,19 @@
 import concurrent.futures as cf
 import sys
 from random import random
-from unittest.mock import patch
+from unittest import mock
 
 import pytest
 
 import sentry_sdk
 from sentry_sdk.integrations.unleash import UnleashIntegration
 from tests.integrations.unleash.testutils import MockUnleashClient
 
-original_is_enabled = MockUnleashClient.is_enabled
-original_get_variant = MockUnleashClient.get_variant
 
-
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_is_enabled(sentry_init, capture_events, uninstall_integration):
     client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration()])
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
 
     client.is_enabled("hello")
     client.is_enabled("world")
@@ -36,11 +32,10 @@ def test_is_enabled(sentry_init, capture_events, uninstall_integration):
     }
 
 
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_get_variant(sentry_init, capture_events, uninstall_integration):
     client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration()])
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
 
     client.get_variant("no_payload_feature")
     client.get_variant("string_feature")
@@ -65,11 +60,10 @@ def test_get_variant(sentry_init, capture_events, uninstall_integration):
     }
 
 
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_is_enabled_threaded(sentry_init, capture_events, uninstall_integration):
     client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration()])
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
     events = capture_events()
 
     def task(flag_key):
@@ -113,11 +107,10 @@ def task(flag_key):
     }
 
 
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_get_variant_threaded(sentry_init, capture_events, uninstall_integration):
     client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration()])
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
     events = capture_events()
 
     def task(flag_key):
@@ -162,13 +155,12 @@ def task(flag_key):
 
 
 @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python3.7 or higher")
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_is_enabled_asyncio(sentry_init, capture_events, uninstall_integration):
     asyncio = pytest.importorskip("asyncio")
 
     client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration()])
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
     events = capture_events()
 
     async def task(flag_key):
@@ -213,13 +205,12 @@ async def runner():
 
 
 @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python3.7 or higher")
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_get_variant_asyncio(sentry_init, capture_events, uninstall_integration):
     asyncio = pytest.importorskip("asyncio")
 
     client = MockUnleashClient()
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration()])
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
     events = capture_events()
 
     async def task(flag_key):
@@ -263,50 +254,41 @@ async def runner():
     }
 
 
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 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
+
     uninstall_integration(UnleashIntegration)
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
+
+    res = client.is_enabled("test-flag", "arg", kwarg=1)
+    assert res == mock_is_enabled.return_value
+    assert mock_is_enabled.call_args == (
+        ("test-flag", "arg"),
+        {"kwarg": 1},
+    )
+
+    res = client.get_variant("test-flag", "arg", kwarg=1)
+    assert res == mock_get_variant.return_value
+    assert mock_get_variant.call_args == (
+        ("test-flag", "arg"),
+        {"kwarg": 1},
+    )
+
 
-    with patch(
-        "sentry_sdk.integrations.unleash.UnleashClient.is_enabled"
-    ) as mock_is_enabled:
-        with patch(
-            "sentry_sdk.integrations.unleash.UnleashClient.get_variant"
-        ) as mock_get_variant:
-            mock_is_enabled.return_value = random() < 0.5
-            mock_get_variant.return_value = {"enabled": random() < 0.5}
-            sentry_init(integrations=[UnleashIntegration()])
-            client = MockUnleashClient()
-
-            res = client.is_enabled("test-flag", "arg", kwarg=1)
-            assert res == mock_is_enabled.return_value
-            assert mock_is_enabled.call_args == (
-                (client, "test-flag", "arg"),
-                {"kwarg": 1},
-            )
-
-            res = client.get_variant("test-flag", "arg", kwarg=1)
-            assert res == mock_get_variant.return_value
-            assert mock_get_variant.call_args == (
-                (client, "test-flag", "arg"),
-                {"kwarg": 1},
-            )
-
-
-@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_wrapper_attributes(sentry_init, uninstall_integration):
+    client = MockUnleashClient()
+    original_is_enabled = client.is_enabled
+    original_get_variant = client.get_variant
+
     uninstall_integration(UnleashIntegration)
-    sentry_init(integrations=[UnleashIntegration()])
+    sentry_init(integrations=[UnleashIntegration(client)])  # type: ignore
 
-    client = MockUnleashClient()
     assert client.is_enabled.__name__ == "is_enabled"
     assert client.is_enabled.__qualname__ == original_is_enabled.__qualname__
-    assert MockUnleashClient.is_enabled.__name__ == "is_enabled"
-    assert MockUnleashClient.is_enabled.__qualname__ == original_is_enabled.__qualname__
 
     assert client.get_variant.__name__ == "get_variant"
     assert client.get_variant.__qualname__ == original_get_variant.__qualname__
-    assert MockUnleashClient.get_variant.__name__ == "get_variant"
-    assert (
-        MockUnleashClient.get_variant.__qualname__ == original_get_variant.__qualname__
-    )