feat: use-env-document-to-reconciliate-types

Author: Zaimwa9

src/edge_proxy/environments.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 from datetime import datetime
 from email.utils import formatdate
 from functools import lru_cache
@@ -83,6 +83,10 @@ def get_flags_response_data(
         server_key_only_feature_ids = environment_document.get("project", {}).get(
             "server_key_only_feature_ids", []
         )
+        feature_types = {
+            fs["feature"]["id"]: fs["feature"].get("type", "STANDARD")
+            for fs in environment_document.get("feature_states", [])
+        }
 
         context = map_environment_document_to_context(environment_document)
         evaluation_result = get_evaluation_result(context)
@@ -100,7 +104,7 @@ def get_flags_response_data(
                 if not filtered:
                     raise FeatureNotFoundError()
 
-            data = map_flag_result_to_response_data(flag_result)
+            data = map_flag_result_to_response_data(flag_result, feature_types)
 
         else:
             flags = list(evaluation_result["flags"].values())
@@ -114,7 +118,7 @@ def get_flags_response_data(
                 "hide_disabled_flags", False
             )
             flags = filter_disabled_flags(flags, hide_disabled_flags)
-            data = map_flag_results_to_response_data(flags)
+            data = map_flag_results_to_response_data(flags, feature_types)
 
         return data
 
@@ -126,6 +130,10 @@ def get_identity_response_data(
         server_key_only_feature_ids = environment_document.get("project", {}).get(
             "server_key_only_feature_ids", []
         )
+        feature_types = {
+            fs["feature"]["id"]: fs["feature"].get("type", "STANDARD")
+            for fs in environment_document.get("feature_states", [])
+        }
 
         environment_context = map_environment_document_to_context(environment_document)
         context = map_context_and_identity_data_to_context(
@@ -147,14 +155,14 @@ def get_identity_response_data(
 
         data = {
             "traits": map_traits_to_response_data(input_data.traits),
-            "flags": map_flag_results_to_response_data(flags),
+            "flags": map_flag_results_to_response_data(flags, feature_types),
         }
         return data
 
     def get_environment(
         self,
         *,
-        environment_key: Optional[str] = None,
+        environment_key: str | None = None,
     ) -> dict[str, Any]:
         if environment_key and environment_key.startswith(SERVER_API_KEY_PREFIX):
             client_side_key = self._get_client_key_from_server_key(environment_key)

src/edge_proxy/feature_utils.py

@@ -5,16 +5,6 @@ def filter_out_server_key_only_flags(
     flags: list[dict[str, Any]],
     server_key_only_feature_ids: list[int],
 ) -> list[dict[str, Any]]:
-    """
-    Filter out server-key-only flags from the list.
-
-    Args:
-        flags: List of flag results
-        server_key_only_feature_ids: List of feature IDs that are server-key-only
-
-    Returns:
-        Filtered list of flags excluding server-key-only features
-    """
     return [
         flag
         for flag in flags
@@ -25,16 +15,6 @@ def filter_out_server_key_only_flags(
 def filter_disabled_flags(
     flags: list[dict[str, Any]], hide_disabled: bool
 ) -> list[dict[str, Any]]:
-    """
-    Filter out disabled flags if hide_disabled is enabled.
-
-    Args:
-        flags: List of flag results
-        hide_disabled: Whether to filter out disabled flags
-
-    Returns:
-        Filtered list of flags excluding disabled ones if hide_disabled is True
-    """
     if not hide_disabled:
         return flags
     return [flag for flag in flags if flag.get("enabled", False)]

src/edge_proxy/mappers.py

@@ -8,13 +8,15 @@
 
 def map_flag_result_to_response_data(
     flag_result: dict[str, Any],
+    feature_types: dict[int, str] | None = None,
 ) -> dict[str, Any]:
-    """Map a single flag result to API response format."""
+    feature_id = flag_result.get("metadata", {}).get("id")
+    feature_type = (feature_types or {}).get(feature_id, "STANDARD")
     return {
         "feature": {
-            "id": flag_result.get("metadata", {}).get("id"),
+            "id": feature_id,
             "name": flag_result["name"],
-            "type": flag_result.get("type", "STANDARD"),
+            "type": feature_type,
         },
         "enabled": flag_result["enabled"],
         "feature_state_value": flag_result["value"],
@@ -23,10 +25,11 @@ def map_flag_result_to_response_data(
 
 def map_flag_results_to_response_data(
     flag_results: list[dict[str, Any]],
+    feature_types: dict[int, str] | None = None,
 ) -> list[dict[str, Any]]:
-    """Map multiple flag results to API response format."""
     return [
-        map_flag_result_to_response_data(flag_result) for flag_result in flag_results
+        map_flag_result_to_response_data(flag_result, feature_types)
+        for flag_result in flag_results
     ]
 
 

tests/test_mappers.py

@@ -75,6 +75,27 @@ def test_map_flag_result_to_response_data__return_expected(
     }
 
 
+def test_map_flag_result_to_response_data__with_feature_types__return_expected(
+    multivariate_flag_result: dict[str, Any],
+) -> None:
+    # Given
+    feature_types = {4: "MULTIVARIATE"}
+
+    # When
+    result = map_flag_result_to_response_data(multivariate_flag_result, feature_types)
+
+    # Then
+    assert result == {
+        "enabled": False,
+        "feature": {
+            "id": 4,
+            "name": "multivariate_feature",
+            "type": "MULTIVARIATE",
+        },
+        "feature_state_value": None,
+    }
+
+
 def test_map_flag_results_to_response_data__return_expected(
     flag_result: dict[str, Any],
     multivariate_flag_result: dict[str, Any],
@@ -108,6 +129,40 @@ def test_map_flag_results_to_response_data__return_expected(
     ]
 
 
+def test_map_flag_results_to_response_data__with_feature_types__return_expected(
+    flag_result: dict[str, Any],
+    multivariate_flag_result: dict[str, Any],
+) -> None:
+    # Given
+    flag_results = [flag_result, multivariate_flag_result]
+    feature_types = {1: "STANDARD", 4: "MULTIVARIATE"}
+
+    # When
+    result = map_flag_results_to_response_data(flag_results, feature_types)
+
+    # Then
+    assert result == [
+        {
+            "enabled": False,
+            "feature": {
+                "id": 1,
+                "name": "feature_1",
+                "type": "STANDARD",
+            },
+            "feature_state_value": "feature_1_value",
+        },
+        {
+            "enabled": False,
+            "feature": {
+                "id": 4,
+                "name": "multivariate_feature",
+                "type": "MULTIVARIATE",
+            },
+            "feature_state_value": None,
+        },
+    ]
+
+
 def test_map_traits_to_response_data__return_expected(
     traits_dict: dict[str, Any],
     traits_data: list[dict[str, Any]],