From f26fc82cd4333a331e7b0da3ec50b02c1bcfc6c8 Mon Sep 17 00:00:00 2001
From: Joe Haines <hello@joehaines.co.uk>
Date: Fri, 25 Aug 2023 10:38:13 +0100
Subject: [PATCH] Use new context module for feature flags

---
 bugsnag/client.py    | 20 ++++----------------
 tests/test_client.py | 15 ++++++---------
 2 files changed, 10 insertions(+), 25 deletions(-)

diff --git a/bugsnag/client.py b/bugsnag/client.py
index 6916c6fb..1fef169c 100644
--- a/bugsnag/client.py
+++ b/bugsnag/client.py
@@ -18,21 +18,11 @@
 from bugsnag.handlers import BugsnagHandler
 from bugsnag.sessiontracker import SessionTracker
 from bugsnag.utils import to_rfc3339
+from bugsnag.context import ContextLocalState
 
 __all__ = ('Client',)
 
 
-try:
-    from contextvars import ContextVar
-    _feature_flag_delegate_context_var = ContextVar(
-        'bugsnag-client-feature-flag-delegate',
-        default=None
-    )  # type: ContextVar[Optional[FeatureFlagDelegate]]
-except ImportError:
-    from bugsnag.utils import ThreadContextVar
-    _feature_flag_delegate_context_var = ThreadContextVar('bugsnag-client-feature-flag-delegate', default=None)  # type: ignore  # noqa: E501
-
-
 class Client:
     """
     A Bugsnag monitoring and reporting client.
@@ -45,6 +35,7 @@ def __init__(self, configuration: Optional[Configuration] = None,
         self.configuration = configuration or Configuration()  # type: Configuration  # noqa: E501
         self.session_tracker = SessionTracker(self.configuration)
         self.configuration.configure(**kwargs)
+        self._context = ContextLocalState(self)
 
         if install_sys_hook:
             self.install_sys_hook()
@@ -237,14 +228,11 @@ def log_handler(
 
     @property
     def _feature_flag_delegate(self) -> FeatureFlagDelegate:
-        try:
-            feature_flag_delegate = _feature_flag_delegate_context_var.get()
-        except LookupError:
-            feature_flag_delegate = None
+        feature_flag_delegate = self._context.feature_flag_delegate
 
         if feature_flag_delegate is None:
             feature_flag_delegate = FeatureFlagDelegate()
-            _feature_flag_delegate_context_var.set(feature_flag_delegate)
+            self._context.feature_flag_delegate = feature_flag_delegate
 
         return feature_flag_delegate
 
diff --git a/tests/test_client.py b/tests/test_client.py
index 0caffcfe..703efe45 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -36,9 +36,6 @@ def setUp(self):
                              asynchronous=False,
                              install_sys_hook=False)
 
-        self.client.clear_feature_flags()
-        legacy.clear_feature_flags()
-
     # Initialisation
 
     def test_init_no_configuration(self):
@@ -1591,16 +1588,16 @@ def test_feature_flags_can_be_added_in_bulk_with_legacy_client(self):
 
     def test_feature_flags_can_be_removed_with_legacy_client(self):
         legacy.add_feature_flags([
-            FeatureFlag('a', '1'),
-            FeatureFlag('b'),
-            FeatureFlag('c', '3')
+            FeatureFlag('x', '1'),
+            FeatureFlag('y'),
+            FeatureFlag('z', '3')
         ])
 
-        legacy.clear_feature_flag('b')
+        legacy.clear_feature_flag('y')
 
         assert legacy.default_client.feature_flags == [
-            FeatureFlag('a', '1'),
-            FeatureFlag('c', '3')
+            FeatureFlag('x', '1'),
+            FeatureFlag('z', '3')
         ]
 
     def test_feature_flags_can_be_cleared_with_legacy_client(self):