From c8104b820d929e64f4ff78a20100d0e8117fee7b Mon Sep 17 00:00:00 2001 From: Andrew Liu <159852527+aliu39@users.noreply.github.com> Date: Mon, 23 Dec 2024 14:07:01 -0800 Subject: [PATCH] Patch an UnleashClient instance instead of the class --- sentry_sdk/integrations/unleash.py | 43 ++++++----- tests/integrations/unleash/test_unleash.py | 86 +++++++++------------- 2 files changed, 60 insertions(+), 69 deletions(-) diff --git a/sentry_sdk/integrations/unleash.py b/sentry_sdk/integrations/unleash.py index 1f90d7ad32..dcad378a15 100644 --- a/sentry_sdk/integrations/unleash.py +++ b/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() diff --git a/tests/integrations/unleash/test_unleash.py b/tests/integrations/unleash/test_unleash.py index 71bcaf1663..98ba37e322 100644 --- a/tests/integrations/unleash/test_unleash.py +++ b/tests/integrations/unleash/test_unleash.py @@ -1,7 +1,7 @@ import concurrent.futures as cf import sys from random import random -from unittest.mock import patch +from unittest import mock import pytest @@ -9,15 +9,11 @@ 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__ - )