wraps_original test

Author: aliu39

tests/integrations/unleash/__init__.py

@@ -1,44 +1,3 @@
 import pytest
 
 pytest.importorskip("UnleashClient")
-
-
-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

@@ -1,22 +1,23 @@
+from random import random
 from unittest.mock import patch
 
 import sentry_sdk
 from sentry_sdk.integrations.unleash import UnleashIntegration
-from tests.integrations.unleash import MockUnleashClient
+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, reset_integration):
-    mock_unleash_client = MockUnleashClient()
+    client = MockUnleashClient()
     reset_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")
+    client.is_enabled("hello")
+    client.is_enabled("world")
+    client.is_enabled("other")
 
     events = capture_events()
     sentry_sdk.capture_exception(Exception("something wrong!"))
@@ -33,16 +34,16 @@ def test_is_enabled(sentry_init, capture_events, reset_integration):
 
 @patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
 def test_get_variant(sentry_init, capture_events, reset_integration):
-    mock_unleash_client = MockUnleashClient()
+    client = MockUnleashClient()
     reset_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")
+    client.get_variant("toggle_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!"))
@@ -56,6 +57,37 @@ def test_get_variant(sentry_init, capture_events, reset_integration):
     }
 
 
+@patch("sentry_sdk.integrations.unleash.UnleashClient", MockUnleashClient)
+def test_wraps_original(sentry_init, reset_integration):
+    reset_integration(UnleashIntegration.identifier)
+
+    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, reset_integration):
     reset_integration(UnleashIntegration.identifier)
     sentry_init(integrations=[UnleashIntegration()])

tests/integrations/unleash/testutils.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)