fix: Return 404 when environment key is not found

src/edge_proxy/environments.py

@@ -14,7 +14,7 @@
 from orjson import orjson
 
 from edge_proxy.cache import BaseEnvironmentsCache, LocalMemEnvironmentsCache
-from edge_proxy.exceptions import FeatureNotFoundError, FlagsmithUnknownKeyError
+from edge_proxy.exceptions import FeatureNotFoundError, UnknownEnvironmentKeyError
 from edge_proxy.feature_utils import filter_out_server_key_only_feature_states
 from edge_proxy.mappers import (
     map_feature_state_to_response_data,
@@ -140,7 +140,7 @@ def get_identity_response_data(
     def get_environment(self, client_side_key: str) -> dict[str, typing.Any]:
         if environment_document := self.cache.get_environment(client_side_key):
             return environment_document
-        raise FlagsmithUnknownKeyError(client_side_key)
+        raise UnknownEnvironmentKeyError(client_side_key)
 
     async def _fetch_document(self, server_side_key: str) -> dict[str, typing.Any]:
         response = await self._client.get(

src/edge_proxy/exceptions.py

@@ -1,4 +1,4 @@
-class FlagsmithUnknownKeyError(KeyError):
+class UnknownEnvironmentKeyError(KeyError):
     pass
 
 

src/edge_proxy/server.py

@@ -5,14 +5,14 @@
 from fastapi import FastAPI, Header
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
-from fastapi.responses import ORJSONResponse
+from fastapi.responses import ORJSONResponse, Response
 
 from edge_proxy.health_check.responses import HealthCheckResponse
 from fastapi_utils.tasks import repeat_every
 
 from edge_proxy.cache import LocalMemEnvironmentsCache
 from edge_proxy.environments import EnvironmentService
-from edge_proxy.exceptions import FeatureNotFoundError, FlagsmithUnknownKeyError
+from edge_proxy.exceptions import FeatureNotFoundError, UnknownEnvironmentKeyError
 from edge_proxy.logging import setup_logging
 from edge_proxy.models import IdentityWithTraits
 from edge_proxy.settings import get_settings
@@ -27,15 +27,9 @@
 app = FastAPI()
 
 
-@app.exception_handler(FlagsmithUnknownKeyError)
+@app.exception_handler(UnknownEnvironmentKeyError)
 async def unknown_key_error(request, exc):
-    return ORJSONResponse(
-        status_code=401,
-        content={
-            "status": "unauthorized",
-            "message": f"unknown key {exc}",
-        },
-    )
+    return Response(status_code=404)
 
 
 @app.get("/health", response_class=ORJSONResponse, deprecated=True)

tests/test_environments.py

@@ -11,7 +11,7 @@
 from edge_proxy.environments import EnvironmentService
 from edge_proxy.exceptions import (
     FeatureNotFoundError,
-    FlagsmithUnknownKeyError,
+    UnknownEnvironmentKeyError,
 )
 from edge_proxy.models import IdentityWithTraits
 from edge_proxy.settings import (
@@ -135,7 +135,7 @@ async def test_get_environment_works_correctly(mocker: MockerFixture):
 
 def test_get_environment_raises_for_unknown_keys():
     environment_service = EnvironmentService(settings=settings)
-    with pytest.raises(FlagsmithUnknownKeyError):
+    with pytest.raises(UnknownEnvironmentKeyError):
         environment_service.get_environment("test_env_key_unknown")
 
 

tests/test_server.py

@@ -156,11 +156,7 @@ def test_get_flags_unknown_key(
         headers={"X-Environment-Key": environment_key},
         params={"feature": "feature_1"},
     )
-    assert response.status_code == 401
-    assert response.json() == {
-        "status": "unauthorized",
-        "message": "unknown key 'unknown_environment_key'",
-    }
+    assert response.status_code == 404
 
 
 def test_post_identity_with_traits(