feat: misc-improvement

Author: Zaimwa9

src/edge_proxy/cache.py

@@ -1,4 +1,4 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 from typing import Any
 
 from edge_proxy.feature_utils import build_feature_types_lookup
@@ -25,18 +25,18 @@ def put_environment(
             return True
         return False
 
+    @abstractmethod
     def _put_environment(
         self,
         environment_api_key: str,
         environment_document: dict[str, Any],
-    ) -> None:
-        raise NotImplementedError()
+    ) -> None: ...
 
-    def get_environment(self, environment_api_key: str) -> dict[str, Any] | None:
-        raise NotImplementedError()
+    @abstractmethod
+    def get_environment(self, environment_api_key: str) -> dict[str, Any] | None: ...
 
-    def get_feature_types(self, environment_api_key: str) -> dict[int, str] | None:
-        return None
+    @abstractmethod
+    def get_feature_types(self, environment_api_key: str) -> dict[int, str] | None: ...
 
 
 _LocalCacheDict = dict[str, dict[str, Any]]

src/edge_proxy/environments.py

@@ -33,6 +33,10 @@
 SERVER_API_KEY_PREFIX = "ser."
 
 
+def _get_hide_disabled_flags(environment_document: dict[str, Any]) -> bool:
+    return environment_document.get("project", {}).get("hide_disabled_flags", False)
+
+
 class EnvironmentService:
     def __init__(
         self,
@@ -77,26 +81,25 @@ async def refresh_environment_caches(self):
             self.last_updated_at = datetime.now()
 
     def get_flags_response_data(
-        self, environment_key: str, feature: str = ""
+        self, environment_key: str, feature: str | None = None
     ) -> 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)
+        if is_server_key:
+            environment_key = self._get_client_key_from_server_key(environment_key)
+
+        environment_document = self.get_environment(environment_key=environment_key)
         server_key_only_feature_ids = environment_document.get("project", {}).get(
             "server_key_only_feature_ids", []
         )
+        hide_disabled_flags = _get_hide_disabled_flags(environment_document)
 
         context = map_environment_document_to_context(environment_document)
         evaluation_result = get_evaluation_result(context)
 
-        feature_types = None
-        if hasattr(self.cache, "get_feature_types"):
-            feature_types = self.cache.get_feature_types(environment_key)
+        feature_types = self.cache.get_feature_types(environment_key)
         if feature_types is None:
             feature_types = build_feature_types_lookup(environment_document)
 
-        data: dict[str, Any] | list[dict[str, Any]]
-
         if feature:
             if feature not in evaluation_result["flags"]:
                 raise FeatureNotFoundError()
@@ -110,38 +113,31 @@ def get_flags_response_data(
                 if not filtered:
                     raise FeatureNotFoundError()
 
-                hide_disabled_flags = environment_document.get("project", {}).get(
-                    "hide_disabled_flags", False
-                )
                 if hide_disabled_flags and not flag_result.get("enabled", False):
                     raise FeatureNotFoundError()
 
-            data = map_flag_result_to_response_data(flag_result, feature_types)
+            return map_flag_result_to_response_data(flag_result, feature_types)
 
-        else:
-            flags = list(evaluation_result["flags"].values())
-
-            if not is_server_key:
-                flags = filter_out_server_key_only_flags(
-                    flags, server_key_only_feature_ids
-                )
-                hide_disabled_flags = environment_document.get("project", {}).get(
-                    "hide_disabled_flags", False
-                )
-                flags = filter_disabled_flags(flags, hide_disabled_flags)
+        flags = list(evaluation_result["flags"].values())
 
-            data = map_flag_results_to_response_data(flags, feature_types)
+        if not is_server_key:
+            flags = filter_out_server_key_only_flags(flags, server_key_only_feature_ids)
+            flags = filter_disabled_flags(flags, hide_disabled_flags)
 
-        return data
+        return map_flag_results_to_response_data(flags, feature_types)
 
     def get_identity_response_data(
         self, input_data: IdentityWithTraits, environment_key: str
     ) -> dict[str, Any]:
-        environment_document = self.get_environment(environment_key=environment_key)
         is_server_key = environment_key.startswith(SERVER_API_KEY_PREFIX)
+        if is_server_key:
+            environment_key = self._get_client_key_from_server_key(environment_key)
+
+        environment_document = self.get_environment(environment_key=environment_key)
         server_key_only_feature_ids = environment_document.get("project", {}).get(
             "server_key_only_feature_ids", []
         )
+        hide_disabled_flags = _get_hide_disabled_flags(environment_document)
 
         environment_context = map_environment_document_to_context(environment_document)
         context = map_context_and_identity_data_to_context(
@@ -151,25 +147,19 @@ def get_identity_response_data(
         )
         evaluation_result = get_evaluation_result(context)
 
-        feature_types = None
-        if hasattr(self.cache, "get_feature_types"):
-            feature_types = self.cache.get_feature_types(environment_key)
+        feature_types = self.cache.get_feature_types(environment_key)
         if feature_types is None:
             feature_types = build_feature_types_lookup(environment_document)
 
         flags = list(evaluation_result["flags"].values())
         if not is_server_key:
             flags = filter_out_server_key_only_flags(flags, server_key_only_feature_ids)
-            hide_disabled_flags = environment_document.get("project", {}).get(
-                "hide_disabled_flags", False
-            )
             flags = filter_disabled_flags(flags, hide_disabled_flags)
 
-        data = {
+        return {
             "traits": map_traits_to_response_data(input_data.traits),
             "flags": map_flag_results_to_response_data(flags, feature_types),
         }
-        return data
 
     def get_environment(
         self,

src/edge_proxy/mappers.py

@@ -1,11 +1,8 @@
 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: FlagResult[Any],
@@ -37,5 +34,4 @@ def map_flag_results_to_response_data(
 def map_traits_to_response_data(
     traits: dict[str, ContextValue],
 ) -> list[dict[str, Any]]:
-    """Convert traits dict to API response format (list of trait objects)."""
     return [{"trait_key": k, "trait_value": v} for k, v in traits.items()]

tests/conftest.py

@@ -70,7 +70,7 @@ def environment_service() -> "EnvironmentService":
 
 
 @pytest.fixture
-def mocked_environment_cache(mocker: MockerFixture):
+def mocked_environment_cache(mocker: MockerFixture) -> typing.Any:
     mock = mocker.patch("edge_proxy.server.environment_service.cache")
     mock.get_environment.return_value = None
     mock.get_feature_types.return_value = None

tests/test_server.py

@@ -316,6 +316,11 @@ def test_get_flags__server_key__hide_disabled_flags_enabled__returns_all_flags(
     assert response.status_code == 200
     flags = response.json()
     assert len(flags) == 3
+    # Verify disabled flags are included (bypasses hide_disabled_flags for server keys)
+    flag_names = {f["feature"]["name"] for f in flags}
+    assert "feature_1" in flag_names  # disabled flag
+    assert "feature_2" in flag_names  # enabled flag
+    assert "feature_3" in flag_names  # disabled flag
 
 
 def test_get_flags__client_key__hide_disabled_flags_enabled__single_disabled_feature__returns_404(