feat: surface the selected variant on flags (#221)

Author: gagantrivedi

flagsmith/api/types.py

@@ -39,6 +39,7 @@ class FeatureSegmentModel(typing.TypedDict):
 
 class MultivariateFeatureOptionModel(typing.TypedDict):
     value: str
+    key: NotRequired[typing.Optional[str]]
 
 
 class MultivariateFeatureStateValueModel(typing.TypedDict):

flagsmith/mappers.py

@@ -7,6 +7,7 @@
 import sseclient
 from flag_engine.context.types import (
     FeatureContext,
+    FeatureValue,
     SegmentContext,
     SegmentRule,
     StrValueSegmentCondition,
@@ -245,11 +246,13 @@ def _map_environment_document_feature_states_to_feature_contexts(
         if multivariate_feature_state_values := feature_state.get(
             "multivariate_feature_state_values"
         ):
-            feature_context["variants"] = [
-                {
-                    "value": multivariate_feature_state_value[
-                        "multivariate_feature_option"
-                    ]["value"],
+            variants: list[FeatureValue] = []
+            for multivariate_feature_state_value in multivariate_feature_state_values:
+                multivariate_feature_option = multivariate_feature_state_value[
+                    "multivariate_feature_option"
+                ]
+                variant: FeatureValue = {
+                    "value": multivariate_feature_option["value"],
                     "weight": multivariate_feature_state_value["percentage_allocation"],
                     "priority": (
                         multivariate_feature_state_value.get("id")
@@ -258,8 +261,10 @@ def _map_environment_document_feature_states_to_feature_contexts(
                         ).int
                     ),
                 }
-                for multivariate_feature_state_value in multivariate_feature_state_values
-            ]
+                if (key := multivariate_feature_option.get("key")) is not None:
+                    variant["key"] = key
+                variants.append(variant)
+            feature_context["variants"] = variants
 
         if feature_segment := feature_state.get("feature_segment"):
             feature_context["priority"] = feature_segment["priority"]

flagsmith/models.py

@@ -51,6 +51,7 @@ class DefaultFlag(BaseFlag):
 class Flag(BaseFlag):
     feature_id: int
     feature_name: str
+    variant: typing.Optional[str] = None
     is_default: bool = field(default=False)
 
     @classmethod
@@ -64,6 +65,7 @@ def from_evaluation_result(
                 value=flag_result["value"],
                 feature_name=flag_result["name"],
                 feature_id=metadata["id"],
+                variant=flag_result.get("variant"),
             )
         raise ValueError(
             "FlagResult metadata is missing. Cannot create Flag instance. "
@@ -77,6 +79,7 @@ def from_api_flag(cls, flag_data: typing.Mapping[str, typing.Any]) -> Flag:
             value=flag_data["feature_state_value"],
             feature_name=flag_data["feature"]["name"],
             feature_id=flag_data["feature"]["id"],
+            variant=flag_data.get("variant"),
         )
 
 

poetry.lock

@@ -10,7 +10,7 @@ documentation = "https://docs.flagsmith.com"
 packages = [{ include = "flagsmith" }]
 
 [tool.poetry.dependencies]
-flagsmith-flag-engine = "^10.0.4"
+flagsmith-flag-engine = "^10.2.0"
 iso8601 = { version = "^2.1.0", python = "<3.11" }
 python = ">=3.9,<4"
 requests = "^2.32.3"

pyproject.toml

@@ -0,0 +1,57 @@
+from flagsmith.api.types import EnvironmentModel
+from flagsmith.mappers import map_environment_document_to_context
+
+
+def _environment_with_keyed_variant() -> EnvironmentModel:
+    return {
+        "api_key": "test-key",
+        "name": "Test",
+        "project": {"segments": []},
+        "identity_overrides": [],
+        "feature_states": [
+            {
+                "enabled": True,
+                "feature": {"id": 1, "name": "mv_feature"},
+                "feature_state_value": "control_value",
+                "featurestate_uuid": "00000000-0000-0000-0000-000000000001",
+                "multivariate_feature_state_values": [
+                    {
+                        "id": 10,
+                        "mv_fs_value_uuid": "00000000-0000-0000-0000-000000000002",
+                        "percentage_allocation": 100,
+                        "multivariate_feature_option": {
+                            "value": "variant_value",
+                            "key": "variant_a",
+                        },
+                    }
+                ],
+            }
+        ],
+    }
+
+
+def test_map_environment_document_to_context__keyed_variant__carries_key() -> None:
+    # Given
+    environment = _environment_with_keyed_variant()
+
+    # When
+    context = map_environment_document_to_context(environment)
+
+    # Then
+    variants = context["features"]["mv_feature"]["variants"]
+    assert variants[0]["key"] == "variant_a"
+
+
+def test_map_environment_document_to_context__null_variant_key__drops_key() -> None:
+    # Given - an unkeyed variant is serialised with a null key
+    environment = _environment_with_keyed_variant()
+    environment["feature_states"][0]["multivariate_feature_state_values"][0][
+        "multivariate_feature_option"
+    ]["key"] = None
+
+    # When
+    context = map_environment_document_to_context(environment)
+
+    # Then - the null key is dropped, treated as no key
+    variants = context["features"]["mv_feature"]["variants"]
+    assert "key" not in variants[0]

tests/test_mappers.py

@@ -22,6 +22,7 @@ def test_flag_from_evaluation_result() -> None:
         "name": "test_feature",
         "reason": "DEFAULT",
         "value": "test-value",
+        "variant": "control",
         "metadata": {"id": 123},
     }
 
@@ -34,6 +35,25 @@ def test_flag_from_evaluation_result() -> None:
     assert flag.feature_name == "test_feature"
     assert flag.feature_id == 123
     assert flag.is_default is False
+    assert flag.variant == "control"
+
+
+def test_flag_from_evaluation_result__no_variant__is_none() -> None:
+    # Given
+    flag_result: SDKFlagResult = {
+        "enabled": True,
+        "name": "test_feature",
+        "reason": "DEFAULT",
+        "value": "test-value",
+        "variant": None,
+        "metadata": {"id": 123},
+    }
+
+    # When
+    flag = Flag.from_evaluation_result(flag_result)
+
+    # Then
+    assert flag.variant is None
 
 
 @pytest.mark.parametrize(
@@ -47,6 +67,7 @@ def test_flag_from_evaluation_result() -> None:
                     "name": "feature1",
                     "reason": "DEFAULT",
                     "value": "value1",
+                    "variant": None,
                     "metadata": {"id": 1},
                 }
             },
@@ -59,6 +80,7 @@ def test_flag_from_evaluation_result() -> None:
                     "name": "feature1",
                     "reason": "DEFAULT",
                     "value": "value1",
+                    "variant": None,
                     "metadata": {"id": 1},
                 }
             },
@@ -71,20 +93,23 @@ def test_flag_from_evaluation_result() -> None:
                     "name": "feature1",
                     "reason": "DEFAULT",
                     "value": "value1",
+                    "variant": None,
                     "metadata": {"id": 1},
                 },
                 "feature2": {
                     "enabled": True,
                     "name": "feature2",
                     "reason": "DEFAULT",
                     "value": "value2",
+                    "variant": None,
                     "metadata": {"id": 2},
                 },
                 "feature3": {
                     "enabled": True,
                     "name": "feature3",
                     "reason": "DEFAULT",
                     "value": 42,
+                    "variant": None,
                     "metadata": {"id": 3},
                 },
             },
@@ -136,6 +161,7 @@ def test_flag_from_evaluation_result_value_types(
         "name": "test_feature",
         "reason": "DEFAULT",
         "value": value,
+        "variant": None,
         "metadata": {"id": 123},
     }
 
@@ -153,13 +179,45 @@ def test_flag_from_evaluation_result_missing_metadata__raises_expected() -> None
         "name": "test_feature",
         "reason": "DEFAULT",
         "value": "test-value",
+        "variant": None,
     }
 
     # When & Then
     with pytest.raises(ValueError):
         Flag.from_evaluation_result(flag_result)
 
 
+def test_flag_from_api_flag__sets_variant() -> None:
+    # Given
+    flag_data = {
+        "enabled": True,
+        "feature_state_value": "test-value",
+        "feature": {"name": "test_feature", "id": 123},
+        "variant": "control",
+    }
+
+    # When
+    flag = Flag.from_api_flag(flag_data)
+
+    # Then
+    assert flag.variant == "control"
+
+
+def test_flag_from_api_flag__no_variant__is_none() -> None:
+    # Given - the REST API may not include `variant`
+    flag_data = {
+        "enabled": True,
+        "feature_state_value": "test-value",
+        "feature": {"name": "test_feature", "id": 123},
+    }
+
+    # When
+    flag = Flag.from_api_flag(flag_data)
+
+    # Then
+    assert flag.variant is None
+
+
 def test_get_flag_without_pipeline_processor() -> None:
     flags = Flags(
         flags={
@@ -197,6 +255,7 @@ def make(
                 "name": "target",
                 "enabled": False,
                 "value": "base-value",
+                "variant": None,
                 "metadata": {"id": 1},
             },
         }
@@ -206,6 +265,7 @@ def make(
                 "name": f"noise_{i}",
                 "enabled": True,
                 "value": f"noise-value-{i}",
+                "variant": None,
                 "metadata": {"id": 100 + i},
             }
         return {