feat: improved-type

Author: Zaimwa9

src/edge_proxy/environments.py

@@ -77,8 +77,8 @@ async def refresh_environment_caches(self):
             self.last_updated_at = datetime.now()
 
     def get_flags_response_data(
-        self, environment_key: str, feature: str = None
-    ) -> dict[str, Any]:
+        self, environment_key: str, feature: str = ""
+    ) -> dict[str, Any] | list[dict[str, Any]]:
         environment_document = self.get_environment(environment_key=environment_key)
         is_server_key = environment_key.startswith(SERVER_API_KEY_PREFIX)
         server_key_only_feature_ids = environment_document.get("project", {}).get(
@@ -92,6 +92,7 @@ def get_flags_response_data(
 
         context = map_environment_document_to_context(environment_document)
         evaluation_result = get_evaluation_result(context)
+        data: dict[str, Any] | list[dict[str, Any]]
 
         if feature:
             if feature not in evaluation_result["flags"]:

src/edge_proxy/feature_utils.py

@@ -1,4 +1,5 @@
 from typing import Any
+from flag_engine.result.types import FlagResult
 
 
 def build_feature_types_lookup(
@@ -11,9 +12,9 @@ def build_feature_types_lookup(
 
 
 def filter_out_server_key_only_flags(
-    flags: list[dict[str, Any]],
+    flags: list[FlagResult[Any]],
     server_key_only_feature_ids: list[int],
-) -> list[dict[str, Any]]:
+) -> list[FlagResult[Any]]:
     return [
         flag
         for flag in flags
@@ -22,8 +23,8 @@ def filter_out_server_key_only_flags(
 
 
 def filter_disabled_flags(
-    flags: list[dict[str, Any]], hide_disabled: bool
-) -> list[dict[str, Any]]:
+    flags: list[FlagResult[Any]], hide_disabled: bool
+) -> list[FlagResult[Any]]:
     if not hide_disabled:
         return flags
     return [flag for flag in flags if flag.get("enabled", False)]

src/edge_proxy/mappers.py

@@ -1,13 +1,14 @@
 from typing import Any
 from flag_engine.engine import ContextValue
+from flag_engine.result.types import FlagResult
 
 from edge_proxy.schemas import APIFeatureStateSchema
 
 _api_feature_state_schema = APIFeatureStateSchema()
 
 
 def map_flag_result_to_response_data(
-    flag_result: dict[str, Any],
+    flag_result: FlagResult[Any],
     feature_types: dict[int, str] | None = None,
 ) -> dict[str, Any]:
     feature_id = flag_result.get("metadata", {}).get("id")
@@ -24,7 +25,7 @@ def map_flag_result_to_response_data(
 
 
 def map_flag_results_to_response_data(
-    flag_results: list[dict[str, Any]],
+    flag_results: list[FlagResult[Any]],
     feature_types: dict[int, str] | None = None,
 ) -> list[dict[str, Any]]:
     return [

tests/fixtures/response_data.py

@@ -50,6 +50,30 @@
 }
 
 
+_multivariate_feature_state = {
+    "multivariate_feature_state_values": [
+        {
+            "id": 1,
+            "multivariate_feature_option": {"id": 1, "value": "variant_a"},
+            "percentage_allocation": 50,
+        },
+        {
+            "id": 2,
+            "multivariate_feature_option": {"id": 2, "value": "variant_b"},
+            "percentage_allocation": 50,
+        },
+    ],
+    "feature_state_value": "control",
+    "featurestate_uuid": "e4g9c5d3-8b52-6g79-1hfd-cg8g54f51i0g",
+    "feature": {
+        "name": "mv_feature",
+        "type": "MULTIVARIATE",
+        "id": 4,
+    },
+    "enabled": True,
+}
+
+
 _segment_1 = {
     "name": "segment_1",
     "rules": [
@@ -157,3 +181,33 @@
     "project": _project_with_hide_disabled_flags,
     "id": 2,
 }
+
+
+_project_with_multivariate = {
+    "name": "project-with-multivariate",
+    "organisation": {
+        "feature_analytics": False,
+        "name": "org-1",
+        "id": 1,
+        "persist_trait_data": True,
+        "stop_serving_flags": False,
+    },
+    "id": 3,
+    "hide_disabled_flags": False,
+    "segments": [],
+    "server_key_only_feature_ids": [],
+}
+
+
+environment_with_multivariate_feature = {
+    "updated_at": "1969-07-20T20:17:40Z",
+    "name": "environment_with_multivariate",
+    "feature_states": [
+        _environment_feature_state_1,
+        _multivariate_feature_state,
+    ],
+    "identity_overrides": [],
+    "api_key": "env_with_multivariate_key",
+    "project": _project_with_multivariate,
+    "id": 3,
+}

tests/test_server.py

@@ -9,6 +9,7 @@
 from tests.fixtures.response_data import (
     environment_1,
     environment_with_hide_disabled_flags,
+    environment_with_multivariate_feature,
 )
 
 if typing.TYPE_CHECKING:
@@ -449,3 +450,36 @@ def test_post_identity__client_key__hide_disabled_flags_disabled__returns_all_fl
     response_data = response.json()
     flags = response_data["flags"]
     assert len(flags) == 2
+
+
+def test_get_flags__multivariate_feature__returns_correct_type(
+    mocker: MockerFixture,
+    client: TestClient,
+) -> None:
+    # Given
+    environment_key = "test_environment_key"
+    mocked_environment_cache = mocker.patch(
+        "edge_proxy.server.environment_service.cache"
+    )
+    mocked_environment_cache.get_environment.return_value = (
+        environment_with_multivariate_feature
+    )
+    mocked_environment_cache.get_feature_types.return_value = None
+
+    # When
+    response = client.get(
+        "/api/v1/flags", headers={"X-Environment-Key": environment_key}
+    )
+
+    # Then
+    assert response.status_code == 200
+    flags = response.json()
+    assert len(flags) == 2
+
+    mv_flag = next(f for f in flags if f["feature"]["name"] == "mv_feature")
+    assert mv_flag["feature"]["type"] == "MULTIVARIATE"
+    assert mv_flag["feature"]["id"] == 4
+    assert mv_flag["enabled"] is True
+
+    standard_flag = next(f for f in flags if f["feature"]["name"] == "feature_1")
+    assert standard_flag["feature"]["type"] == "STANDARD"