fix: get_identity_segments tries to return identity override segments

flagsmith/flagsmith.py

@@ -14,6 +14,7 @@
     map_context_and_identity_data_to_context,
     map_environment_document_to_context,
     map_environment_document_to_environment_updated_at,
+    map_segment_results_to_identity_segments,
 )
 from flagsmith.models import DefaultFlag, Flags, Segment
 from flagsmith.offline_handlers import OfflineHandler
@@ -283,10 +284,8 @@ def get_identity_segments(
         evaluation_result = engine.get_evaluation_result(
             context=context,
         )
-        return [
-            Segment(id=int(segment_result["key"]), name=segment_result["name"])
-            for segment_result in evaluation_result["segments"]
-        ]
+
+        return map_segment_results_to_identity_segments(evaluation_result["segments"])
 
     def update_environment(self) -> None:
         try:

flagsmith/mappers.py

@@ -11,9 +11,11 @@
     SegmentContext,
     SegmentRule,
 )
+from flag_engine.result.types import SegmentResult
 from flag_engine.segments.types import ContextValue
 
-from flagsmith.types import StreamEvent, TraitConfig
+from flagsmith.models import Segment
+from flagsmith.types import SegmentMetadata, StreamEvent, TraitConfig
 
 OverrideKey = typing.Tuple[
     str,
@@ -24,6 +26,25 @@
 OverridesKey = typing.Tuple[OverrideKey, ...]
 
 
+def map_segment_results_to_identity_segments(
+    segment_results: list[SegmentResult],
+) -> list[Segment]:
+    identity_segments: list[Segment] = []
+    for segment_result in segment_results:
+        if raw_metadata := segment_result.get("metadata"):
+            metadata = typing.cast(SegmentMetadata, raw_metadata)
+            if metadata.get("source") == "api" and (
+                (flagsmith_id := metadata.get("flagsmith_id")) is not None
+            ):
+                identity_segments.append(
+                    Segment(
+                        id=flagsmith_id,
+                        name=segment_result["name"],
+                    )
+                )
+    return identity_segments
+
+
 def map_sse_event_to_stream_event(event: sseclient.Event) -> StreamEvent:
     event_data = json.loads(event.data)
     return {
@@ -90,7 +111,7 @@ def map_environment_document_to_context(
         },
         "segments": {
             **{
-                (segment_key := str(segment["id"])): {
+                (segment_key := str(segment_id := segment["id"])): {
                     "key": segment_key,
                     "name": segment["name"],
                     "rules": _map_environment_document_rules_to_context_rules(
@@ -101,6 +122,12 @@ def map_environment_document_to_context(
                             segment.get("feature_states") or []
                         )
                     ),
+                    "metadata": dict(
+                        SegmentMetadata(
+                            flagsmith_id=segment_id,
+                            source="api",
+                        )
+                    ),
                 }
                 for segment in environment_document["project"]["segments"]
             },
@@ -142,6 +169,7 @@ def _map_identity_overrides_to_segments(
         # Create a segment context for each unique set of overrides
         # Generate a unique key to avoid collisions
         segment_key = str(hash(overrides_key))
+        segment_metadata = SegmentMetadata(source="identity_overrides")
         segment_contexts[segment_key] = SegmentContext(
             key="",  # Identity override segments never use % Split operator
             name="identity_overrides",
@@ -168,6 +196,7 @@ def _map_identity_overrides_to_segments(
                 }
                 for feature_key, feature_name, feature_enabled, feature_value in overrides_key
             ],
+            metadata=dict(segment_metadata),
         )
     return segment_contexts
 

flagsmith/models.py

@@ -63,13 +63,13 @@ def from_evaluation_result(
     ) -> Flags:
         return cls(
             flags={
-                flag["name"]: Flag(
+                flag_name: Flag(
                     enabled=flag["enabled"],
                     value=flag["value"],
                     feature_name=flag["name"],
                     feature_id=int(flag["feature_key"]),
                 )
-                for flag in evaluation_result["flags"]
+                for flag_name, flag in evaluation_result["flags"].items()
             },
             default_flag_handler=default_flag_handler,
             _analytics_processor=analytics_processor,

flagsmith/types.py

@@ -33,3 +33,10 @@ class TraitConfig(typing.TypedDict):
 class ApplicationMetadata(typing.TypedDict):
     name: NotRequired[str]
     version: NotRequired[str]
+
+
+class SegmentMetadata(typing.TypedDict):
+    flagsmith_id: NotRequired[int]
+    """The ID of the segment used in Flagsmith API."""
+    source: NotRequired[typing.Literal["api", "identity_overrides"]]
+    """The source of the segment, e.g. 'api', 'identity_overrides'."""

pyproject.toml

@@ -13,7 +13,7 @@ packages = [{ include = "flagsmith" }]
 python = ">=3.9,<4"
 requests = "^2.32.3"
 requests-futures = "^1.0.1"
-flagsmith-flag-engine = "^7.0.0"
+flagsmith-flag-engine = "^8.0.0"
 sseclient-py = "^1.8.0"
 
 [tool.poetry.group.dev]

tests/test_flagsmith.py

@@ -159,14 +159,14 @@ def test_get_identity_flags_uses_local_environment_when_available(
     mock_engine = mocker.patch("flagsmith.flagsmith.engine")
 
     expected_evaluation_result = {
-        "flags": [
-            {
+        "flags": {
+            "some_feature": {
                 "name": "some_feature",
                 "enabled": True,
                 "value": "some-feature-state-value",
                 "feature_key": "1",
             }
-        ],
+        },
         "segments": [],
     }
 
@@ -509,6 +509,26 @@ def test_get_identity_segments_with_valid_trait(
     assert segments[0].name == "Test segment"  # obtained from data/environment.json
 
 
+def test_get_identity_segments__identity_overrides__returns_expected(
+    local_eval_flagsmith: Flagsmith,
+) -> None:
+    # Given
+    # the identifier matches the identity override in data/environment.json
+    identifier = "overridden-id"
+    # traits match the "Test segment" segment in data/environment.json
+    traits = {"foo": "bar"}
+
+    # When
+    segments = local_eval_flagsmith.get_identity_segments(identifier, traits)
+
+    # Then
+    # identity override virtual segment is not returned,
+    # only the segment matching the traits
+    assert len(segments) == 1
+    assert segments[0].id == 1
+    assert segments[0].name == "Test segment"
+
+
 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_models.py

@@ -29,55 +29,66 @@ def test_flag_from_evaluation_result() -> None:
 @pytest.mark.parametrize(
     "flags_result,expected_count,expected_names",
     [
-        ([], 0, []),
+        ({}, 0, []),
         (
-            [
-                {
+            {
+                "feature1": {
                     "name": "feature1",
                     "enabled": True,
                     "value": "value1",
                     "feature_key": "1",
                 }
-            ],
+            },
             1,
             ["feature1"],
         ),
         (
-            [
-                {
+            {
+                "feature1": {
+                    "name": "feature1",
+                    "enabled": True,
+                    "value": "value1",
+                    "feature_key": "1",
+                }
+            },
+            1,
+            ["feature1"],
+        ),
+        (
+            {
+                "feature1": {
                     "name": "feature1",
                     "enabled": True,
                     "value": "value1",
                     "feature_key": "1",
                 },
-                {
+                "feature2": {
                     "name": "feature2",
-                    "enabled": False,
-                    "value": None,
+                    "enabled": True,
+                    "value": "value2",
                     "feature_key": "2",
                 },
-                {"name": "feature3", "enabled": True, "value": 42, "feature_key": "3"},
-            ],
+                "feature3": {
+                    "name": "feature3",
+                    "enabled": True,
+                    "value": 42,
+                    "feature_key": "3",
+                },
+            },
             3,
             ["feature1", "feature2", "feature3"],
         ),
     ],
 )
 def test_flags_from_evaluation_result(
-    flags_result: typing.List[FlagResult],
+    flags_result: typing.Dict[str, FlagResult],
     expected_count: int,
     expected_names: typing.List[str],
 ) -> None:
     # Given
     evaluation_result: EvaluationResult = {
         "flags": flags_result,
         "segments": [],
-        "context": {
-            "environment": {
-                "name": "test_environment",
-                "key": "test_environment_key",
-            }
-        },
     }
 
     # When