fix: Identity overrides silently lost during local evaluation (#240)

Author: khvn26

flagsmith/mappers.py

@@ -177,7 +177,7 @@ def _map_identity_overrides_to_segments(
         # Generate a unique key to avoid collisions
         segment_key = str(hash(overrides_key))
         segment_contexts[segment_key] = SegmentContext(
-            key="",  # Identity override segments never use % Split operator
+            key=segment_key,
             name="identity_overrides",
             rules=[
                 {

flagsmith/models.py

@@ -17,7 +17,7 @@
 )
 
 SegmentOverridesIndex = typing.Dict[
-    str, typing.List[SegmentContext[SegmentMetadata, FeatureMetadata]]
+    str, typing.Dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]]
 ]
 
 
@@ -28,11 +28,16 @@ def build_segment_overrides_index(
 
     Computed once per environment-document refresh so the lazy eval path
     can walk only the segments actually relevant to a given flag.
+
+    Each entry keeps the context's own segment keys, so the trimmed
+    contexts built by `Flags._resolve_flag` stay a faithful subset of
+    the environment. Re-keying them by anything else risks collisions
+    silently dropping segments.
     """
     index: SegmentOverridesIndex = {}
-    for segment_context in (context.get("segments") or {}).values():
+    for segment_key, segment_context in (context.get("segments") or {}).items():
         for override in segment_context.get("overrides") or ():
-            index.setdefault(override["name"], []).append(segment_context)
+            index.setdefault(override["name"], {})[segment_key] = segment_context
     return index
 
 
@@ -254,10 +259,7 @@ def _resolve_flag(self, feature_name: str) -> Flag:
         trimmed: SDKEvaluationContext = {
             **context,
             "features": {feature_name: context["features"][feature_name]},
-            "segments": {
-                segment_context["key"]: segment_context
-                for segment_context in overrides_index.get(feature_name, ())
-            },
+            "segments": overrides_index.get(feature_name, {}),
         }
         result = engine.get_evaluation_result(trimmed)
         return Flag.from_evaluation_result(result["flags"][feature_name])

tests/data/environment.json

@@ -79,6 +79,29 @@
                     "feature_segment": null
                 }
             ]
+        },
+        {
+            "identifier": "another-overridden-id",
+            "identity_uuid": "3a2f0d1e-4b6c-4a58-9a1d-2c5f8e7b6d40",
+            "created_date": "2019-08-27T14:53:45.698555Z",
+            "updated_at": "2023-07-14T16:12:00.000000Z",
+            "environment_api_key": "B62qaMZNwfiqT76p38ggrQ",
+            "identity_features": [
+                {
+                    "id": 2,
+                    "feature": {
+                        "id": 1,
+                        "name": "some_feature",
+                        "type": "STANDARD"
+                    },
+                    "featurestate_uuid": "c1f6ba18-9c2d-4a3e-8f57-1d3b9e0a24c7",
+                    "feature_state_value": "another-overridden-value",
+                    "enabled": false,
+                    "environment": 1,
+                    "identity": null,
+                    "feature_segment": null
+                }
+            ]
         }
     ]
 }

tests/test_flagsmith.py

@@ -642,6 +642,22 @@ def test_get_identity_segments__identity_overrides__returns_expected(
     assert segments[0].name == "Test segment"
 
 
+def test_get_identity_flags__local_evaluation_distinct_identity_overrides__returns_expected(
+    local_eval_flagsmith: Flagsmith,
+) -> None:
+    # Given / When
+    first_flag = local_eval_flagsmith.get_identity_flags("overridden-id").get_flag(
+        "some_feature"
+    )
+    second_flag = local_eval_flagsmith.get_identity_flags(
+        "another-overridden-id"
+    ).get_flag("some_feature")
+
+    # Then
+    assert first_flag.value == "some-overridden-value"
+    assert second_flag.value == "another-overridden-value"
+
+
 def test_local_evaluation_requires_server_key() -> None:
     with pytest.raises(ValueError):
         Flagsmith(environment_key="not-a-server-key", enable_local_evaluation=True)

tests/test_mappers.py

@@ -55,3 +55,23 @@ def test_map_environment_document_to_context__null_variant_key__drops_key() -> N
     # Then - the null key is dropped, treated as no key
     variants = context["features"]["mv_feature"]["variants"]
     assert "key" not in variants[0]
+
+
+def test_map_environment_document_to_context__identity_overrides__keys_are_unique(
+    environment: EnvironmentModel,
+) -> None:
+    # Given / When
+    context = map_environment_document_to_context(environment)
+
+    # Then
+    segments = context["segments"] or {}
+    identity_override_keys = {
+        segment_key
+        for segment_key, segment_context in segments.items()
+        if segment_context["name"] == "identity_overrides"
+    }
+    assert len(identity_override_keys) == 2
+    assert all(
+        segments[segment_key]["key"] == segment_key
+        for segment_key in identity_override_keys
+    )

tests/test_models.py

@@ -466,7 +466,8 @@ def default(name: str) -> DefaultFlag:
 def test_build_segment_overrides_index__indexes_only_overriding_segments(
     lazy_context: SDKEvaluationContext,
 ) -> None:
-    # Given: a second segment with no overrides on top of the default context.
+    # Given
+    # a second segment with no overrides on top of the default context
     assert lazy_context["segments"] is not None
     lazy_context["segments"]["no_override_segment"] = {
         "key": "no_override_segment",
@@ -481,9 +482,32 @@ def test_build_segment_overrides_index__indexes_only_overriding_segments(
         ],
     }
 
-    # When: we build the reverse index.
+    # When
+    # we build the reverse index
     index = build_segment_overrides_index(lazy_context)
 
-    # Then: only segments that actually carry an override appear.
+    # Then
+    # only segments that actually carry an override appear,
+    # keyed by their key in the evaluation context
     assert set(index) == {"target"}
-    assert index["target"][0]["name"] == "premium_segment"
+    assert index["target"]["premium_segment"]["name"] == "premium_segment"
+
+
+def test_build_segment_overrides_index__duplicate_segment_key_fields__keeps_both(
+    lazy_context: SDKEvaluationContext,
+) -> None:
+    # Given: a second overriding segment reusing the first segment's `key`
+    # field, as identity-override segments used to do.
+    assert lazy_context["segments"] is not None
+    premium_segment = lazy_context["segments"]["premium_segment"]
+    lazy_context["segments"]["enterprise_segment"] = {
+        **premium_segment,
+        "name": "enterprise_segment",
+    }
+
+    # When
+    index = build_segment_overrides_index(lazy_context)
+
+    # Then
+    # both segments are indexed, neither collapses into the other
+    assert set(index["target"]) == {"premium_segment", "enterprise_segment"}