From 4cc1cae67cd31edbf84ea6d46fcc79435f8609fb Mon Sep 17 00:00:00 2001
From: Andrew Liu <159852527+aliu39@users.noreply.github.com>
Date: Sun, 22 Dec 2024 20:49:10 -0800
Subject: [PATCH] Track all variants regardless of payload

---
 sentry_sdk/integrations/unleash.py         | 15 ++++++++-------
 tests/integrations/unleash/test_unleash.py | 20 ++++++++++++--------
 tests/integrations/unleash/testutils.py    |  2 +-
 3 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/sentry_sdk/integrations/unleash.py b/sentry_sdk/integrations/unleash.py
index 98ec51939f..1f90d7ad32 100644
--- a/sentry_sdk/integrations/unleash.py
+++ b/sentry_sdk/integrations/unleash.py
@@ -30,8 +30,8 @@ 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..
+            # We have no way of knowing what type of unleash feature this is, so we have to treat
+            # it as a boolean / toggle feature.
             flags = sentry_sdk.get_current_scope().flags
             flags.set(feature, enabled)
 
@@ -42,12 +42,13 @@ 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")
-
-            if payload_type is None:
-                flags = sentry_sdk.get_current_scope().flags
-                flags.set(feature, enabled)
+            # _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
+            # variant as a boolean toggle, using the `enabled` field.
+            flags = sentry_sdk.get_current_scope().flags
+            flags.set(feature, enabled)
             return variant
 
         UnleashClient.is_enabled = sentry_is_enabled  # type: ignore
diff --git a/tests/integrations/unleash/test_unleash.py b/tests/integrations/unleash/test_unleash.py
index 2dc49e5a43..71bcaf1663 100644
--- a/tests/integrations/unleash/test_unleash.py
+++ b/tests/integrations/unleash/test_unleash.py
@@ -42,7 +42,7 @@ def test_get_variant(sentry_init, capture_events, uninstall_integration):
     uninstall_integration(UnleashIntegration)
     sentry_init(integrations=[UnleashIntegration()])
 
-    client.get_variant("toggle_feature")
+    client.get_variant("no_payload_feature")
     client.get_variant("string_feature")
     client.get_variant("json_feature")
     client.get_variant("csv_feature")
@@ -55,7 +55,11 @@ def test_get_variant(sentry_init, capture_events, uninstall_integration):
     assert len(events) == 1
     assert events[0]["contexts"]["flags"] == {
         "values": [
-            {"flag": "toggle_feature", "result": True},
+            {"flag": "no_payload_feature", "result": True},
+            {"flag": "string_feature", "result": True},
+            {"flag": "json_feature", "result": True},
+            {"flag": "csv_feature", "result": True},
+            {"flag": "number_feature", "result": True},
             {"flag": "unknown_feature", "result": False},
         ]
     }
@@ -129,7 +133,7 @@ def task(flag_key):
     client.get_variant("hello")
 
     with cf.ThreadPoolExecutor(max_workers=2) as pool:
-        pool.map(task, ["other", "toggle_feature"])
+        pool.map(task, ["no_payload_feature", "other"])
 
     # Capture error in original scope
     sentry_sdk.set_tag("task_id", "0")
@@ -146,13 +150,13 @@ def task(flag_key):
     assert events[1]["contexts"]["flags"] == {
         "values": [
             {"flag": "hello", "result": False},
-            {"flag": "other", "result": False},
+            {"flag": "no_payload_feature", "result": True},
         ]
     }
     assert events[2]["contexts"]["flags"] == {
         "values": [
             {"flag": "hello", "result": False},
-            {"flag": "toggle_feature", "result": True},
+            {"flag": "other", "result": False},
         ]
     }
 
@@ -226,7 +230,7 @@ async def task(flag_key):
             sentry_sdk.capture_exception(Exception("something wrong!"))
 
     async def runner():
-        return asyncio.gather(task("other"), task("toggle_feature"))
+        return asyncio.gather(task("no_payload_feature"), task("other"))
 
     # Capture an eval before we split isolation scopes.
     client.get_variant("hello")
@@ -248,13 +252,13 @@ async def runner():
     assert events[1]["contexts"]["flags"] == {
         "values": [
             {"flag": "hello", "result": False},
-            {"flag": "other", "result": False},
+            {"flag": "no_payload_feature", "result": True},
         ]
     }
     assert events[2]["contexts"]["flags"] == {
         "values": [
             {"flag": "hello", "result": False},
-            {"flag": "toggle_feature", "result": True},
+            {"flag": "other", "result": False},
         ]
     }
 
diff --git a/tests/integrations/unleash/testutils.py b/tests/integrations/unleash/testutils.py
index 582105f3f3..1f6f08fca8 100644
--- a/tests/integrations/unleash/testutils.py
+++ b/tests/integrations/unleash/testutils.py
@@ -27,7 +27,7 @@ def __init__(self, *a, **kw):
                 "enabled": True,
                 "payload": {"type": "csv", "value": "abc 123\ncsbq 94"},
             },
-            "toggle_feature": {"name": "variant1", "enabled": True},
+            "no_payload_feature": {"name": "variant1", "enabled": True},
         }
 
         self.disabled_variant = {"name": "disabled", "enabled": False}