From 5bfb3eceb964042e66624a68dc22cc037f24b18e Mon Sep 17 00:00:00 2001
From: ash-vd <contact@ashwinvandijk.nl>
Date: Wed, 14 Oct 2026 08:42:28 +0200
Subject: [PATCH] fix: return 401 instead of 500 when X-Environment-Key is
 missing

Since #182, get_flags_response_data and get_identity_response_data call
environment_key.startswith() before get_environment(), so a request
without an X-Environment-Key header raises AttributeError and returns a
500. Up to 2.21.2 the same request got a 401 from get_environment().

Check for a missing key first and raise FlagsmithUnknownKeyError, the
same response the environment document endpoint gives.
---
 src/edge_proxy/environments.py | 10 ++++++++--
 tests/test_server.py           | 23 +++++++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/edge_proxy/environments.py b/src/edge_proxy/environments.py
index 880930e..4b144fb 100644
--- a/src/edge_proxy/environments.py
+++ b/src/edge_proxy/environments.py
@@ -82,8 +82,11 @@ async def refresh_environment_caches(self):
             self.last_updated_at = datetime.now()
 
     def get_flags_response_data(
-        self, environment_key: str, feature: str | None = None
+        self, environment_key: str | None, feature: str | None = None
     ) -> dict[str, Any] | list[dict[str, Any]]:
+        if not environment_key:
+            raise FlagsmithUnknownKeyError(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)
@@ -128,8 +131,11 @@ def get_flags_response_data(
         return map_flag_results_to_response_data(flags, feature_types)
 
     def get_identity_response_data(
-        self, input_data: IdentityWithTraits, environment_key: str
+        self, input_data: IdentityWithTraits, environment_key: str | None
     ) -> dict[str, Any]:
+        if not environment_key:
+            raise FlagsmithUnknownKeyError(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)
diff --git a/tests/test_server.py b/tests/test_server.py
index e678d6b..badad2d 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -87,6 +87,16 @@ def test_get_flags_unknown_key(
     }
 
 
+def test_get_flags_missing_key(
+    client: TestClient,
+) -> None:
+    # When
+    response = client.get("/api/v1/flags")
+
+    # Then
+    assert response.status_code == 401
+
+
 def test_post_identity_with_traits(
     mocked_environment_cache,
     environment_1_feature_states_response_list_response_with_segment_override,
@@ -171,6 +181,19 @@ def test_post_identity__invalid_trait_data__expected_response(
     assert response.json()["detail"][-1]["type"] == "string_too_long"
 
 
+def test_post_identity_missing_key(
+    client: TestClient,
+) -> None:
+    # When
+    response = client.post(
+        "/api/v1/identities/",
+        content=orjson.dumps({"identifier": "test_identifier"}),
+    )
+
+    # Then
+    assert response.status_code == 401
+
+
 def test_get_identities(
     mocked_environment_cache,
     client: TestClient,