Populate feature_flags_by_key when setting feature_flags

Since `self.feature_flags_by_key` is derived from `self.feature_flags`, and we often set the latter in unit tests, but forget to set the former, our tests can be wonky.

This ensures that when we set `self.feature_flags`, we always set `self. feature_flags_by_key`

Author: haacked

posthog/client.py

@@ -135,7 +135,7 @@ def __init__(
         self.host = determine_server_host(host)
         self.gzip = gzip
         self.timeout = timeout
-        self.feature_flags = None
+        self._feature_flags = None  # private variable to store flags
         self.feature_flags_by_key = None
         self.group_type_mapping = None
         self.cohorts = None
@@ -204,6 +204,17 @@ def __init__(
                 if send:
                     consumer.start()
 
+    @property
+    def feature_flags(self):
+        return self._feature_flags
+
+    @feature_flags.setter
+    def feature_flags(self, flags):
+        self._feature_flags = flags or []
+        self.feature_flags_by_key = {
+            flag["key"]: flag for flag in self._feature_flags if flag.get("key") is not None
+        }
+
     def identify(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
         if context is not None:
             warnings.warn(
@@ -640,9 +651,6 @@ def _load_feature_flags(self):
             )
 
             self.feature_flags = response["flags"] or []
-            self.feature_flags_by_key = {
-                flag["key"]: flag for flag in self.feature_flags if flag.get("key") is not None
-            }
             self.group_type_mapping = response["group_type_mapping"] or {}
             self.cohorts = response["cohorts"] or {}
 
@@ -664,7 +672,6 @@ def _load_feature_flags(self):
                 )
                 # Reset all feature flag data when quota limited
                 self.feature_flags = []
-                self.feature_flags_by_key = {}
                 self.group_type_mapping = {}
                 self.cohorts = {}
 

posthog/test/test_feature_flags.py

@@ -768,7 +768,6 @@ def test_get_all_flags_and_payloads_with_no_fallback(self, patch_decide, patch_c
             basic_flag,
             disabled_flag,
         ]
-        client.feature_flags_by_key = {"beta-feature": basic_flag, "disabled-feature": disabled_flag}
         all_flags_and_payloads = client.get_all_flags_and_payloads("distinct_id")
         self.assertEqual(
             client.to_payloads(all_flags_and_payloads), {"beta-feature": "new"}
@@ -900,7 +899,6 @@ def test_get_all_flags_and_payloads_with_fallback_but_only_local_evaluation_set(
             flag_2,
             flag_3,
         ]
-        client.feature_flags_by_key = {"beta-feature": flag_1, "disabled-feature": flag_2, "beta-feature2": flag_3}
         # beta-feature2 has no value
         flags_and_payloads = client.get_all_flags_and_payloads("distinct_id", only_evaluate_locally=True)
         self.assertEqual(
@@ -1629,8 +1627,7 @@ def test_boolean_feature_flag_payloads_local(self, patch_decide):
             },
         }
         self.client.feature_flags = [basic_flag]
-        self.client.feature_flags_by_key = {"person-flag": basic_flag}
-
+        
         self.assertEqual(
             self.client.get_feature_flag_payload(
                 "person-flag", "some-distinct-id", person_properties={"region": "USA"}
@@ -1697,8 +1694,7 @@ def test_multivariate_feature_flag_payloads(self, patch_decide):
             },
         }
         self.client.feature_flags = [multivariate_flag]
-        self.client.feature_flags_by_key = {"beta-feature": multivariate_flag}
-
+        
         self.assertEqual(
             self.client.get_feature_flag_payload(
                 "beta-feature", "test_id", person_properties={"email": "[email protected]"}