diff --git a/src/edge_proxy/environments.py b/src/edge_proxy/environments.py index d3b6c7e..29c313c 100644 --- a/src/edge_proxy/environments.py +++ b/src/edge_proxy/environments.py @@ -1,4 +1,4 @@ -import typing +from typing import Any, Optional from datetime import datetime from email.utils import formatdate from functools import lru_cache @@ -76,8 +76,8 @@ async def refresh_environment_caches(self): def get_flags_response_data( self, environment_key: str, feature: str = None - ) -> dict[str, typing.Any]: - environment_document = self.get_environment(environment_key) + ) -> dict[str, Any]: + environment_document = self.get_environment(environment_key=environment_key) environment = EnvironmentModel.model_validate(environment_document) is_server_key = environment_key.startswith(SERVER_API_KEY_PREFIX) @@ -105,8 +105,8 @@ def get_flags_response_data( def get_identity_response_data( self, input_data: IdentityWithTraits, environment_key: str - ) -> dict[str, typing.Any]: - environment_document = self.get_environment(environment_key) + ) -> dict[str, Any]: + environment_document = self.get_environment(environment_key=environment_key) environment = EnvironmentModel.model_validate(environment_document) is_server_key = environment_key.startswith(SERVER_API_KEY_PREFIX) @@ -137,14 +137,22 @@ def get_identity_response_data( } return data - def get_environment(self, client_side_key: str) -> dict[str, typing.Any]: + def get_environment( + self, + *, + environment_key: Optional[str] = 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) + else: + client_side_key = environment_key + if environment_document := self.cache.get_environment(client_side_key): return environment_document - raise FlagsmithUnknownKeyError(client_side_key) - async def _fetch_document( - self, key_pair: EnvironmentKeyPair - ) -> dict[str, typing.Any]: + raise FlagsmithUnknownKeyError(environment_key) + + async def _fetch_document(self, key_pair: EnvironmentKeyPair) -> dict[str, Any]: headers = { "X-Environment-Key": key_pair.server_side_key, } @@ -186,3 +194,9 @@ async def _clear_endpoint_caches(self): func.cache_clear() except AttributeError: pass + + def _get_client_key_from_server_key(self, server_key: str) -> str: + for key_pair in self.settings.environment_key_pairs: + if key_pair.server_side_key == server_key: + return key_pair.client_side_key + raise FlagsmithUnknownKeyError(server_key) diff --git a/src/edge_proxy/server.py b/src/edge_proxy/server.py index 95e6836..372109d 100644 --- a/src/edge_proxy/server.py +++ b/src/edge_proxy/server.py @@ -125,6 +125,17 @@ async def get_identities( return ORJSONResponse(data) +@app.get("/api/v1/environment-document", response_class=ORJSONResponse) +async def environment_document( + x_environment_key: str = Header(None), +) -> ORJSONResponse: + if environment_doc := environment_service.get_environment( + environment_key=x_environment_key, + ): + return ORJSONResponse(environment_doc) + return ORJSONResponse(status_code=401, content=None) + + app.add_middleware( CORSMiddleware, allow_origins=settings.allow_origins, diff --git a/tests/test_environments.py b/tests/test_environments.py index edc55c2..3609f4d 100644 --- a/tests/test_environments.py +++ b/tests/test_environments.py @@ -111,13 +111,13 @@ async def test_get_environment_works_correctly(mocker: MockerFixture): # Next, test that get environment return correct document assert ( environment_service.get_environment( - settings.environment_key_pairs[0].client_side_key + environment_key=settings.environment_key_pairs[0].client_side_key ) == doc_1 ) assert ( environment_service.get_environment( - settings.environment_key_pairs[1].client_side_key + environment_key=settings.environment_key_pairs[1].client_side_key ) == doc_2 ) @@ -125,10 +125,10 @@ async def test_get_environment_works_correctly(mocker: MockerFixture): # Next, let's verify that any additional call to get_environment does not call fetch document environment_service.get_environment( - settings.environment_key_pairs[0].client_side_key + environment_key=settings.environment_key_pairs[0].client_side_key ) environment_service.get_environment( - settings.environment_key_pairs[1].client_side_key + environment_key=settings.environment_key_pairs[1].client_side_key ) assert mock_client.get.call_count == 2 @@ -136,7 +136,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): - environment_service.get_environment("test_env_key_unknown") + environment_service.get_environment(environment_key="test_env_key_unknown") @pytest.mark.asyncio @@ -278,9 +278,7 @@ async def test_get_flags_response_data_skips_filter_for_server_key( # We create a new settings object that contains a server key as a client_side_key api_key = "ser." + environment_1_api_key _settings = AppSettings( - environment_key_pairs=[ - {"client_side_key": api_key, "server_side_key": "ser.key"} - ] + environment_key_pairs=[{"client_side_key": api_key, "server_side_key": api_key}] ) mocked_client = mocker.AsyncMock() @@ -342,9 +340,7 @@ async def test_get_identity_flags_response_skips_filter_for_server_key( # We create a new settings object that contains a server key as a client_side_key api_key = "ser." + environment_1_api_key _settings = AppSettings( - environment_key_pairs=[ - {"client_side_key": api_key, "server_side_key": "ser.key"} - ] + environment_key_pairs=[{"client_side_key": api_key, "server_side_key": api_key}] ) mocked_client = mocker.AsyncMock() diff --git a/tests/test_server.py b/tests/test_server.py index 55c4c0b..6d4c539 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -5,6 +5,7 @@ from pytest_mock import MockerFixture from edge_proxy.main import serve +from edge_proxy.settings import EnvironmentKeyPair from tests.fixtures.response_data import environment_1 if typing.TYPE_CHECKING: @@ -233,3 +234,51 @@ def test_serve_passes_proxy_headers_setting(mocker: MockerFixture) -> None: # Then _, kwargs = mock_uvicorn.call_args assert kwargs.get("proxy_headers") is True + + +def test_get_environment_document( + mocker: MockerFixture, + client: TestClient, +) -> None: + # Given + environment_key_pairs = [ + EnvironmentKeyPair(server_side_key="ser.good", client_side_key="foo") + ] + mocker.patch( + "edge_proxy.server.settings.environment_key_pairs", environment_key_pairs + ) + mocker.patch( + "edge_proxy.server.environment_service.cache" + ).get_environment.return_value = environment_1 + + # When + response = client.get( + "/api/v1/environment-document", + headers={"X-Environment-Key": environment_key_pairs[0].server_side_key}, + ) + + # Then + assert response.status_code == 200 + assert response.json() == environment_1 + + +def test_get_environment_document_missing_key( + client: TestClient, +) -> None: + # When + response = client.get( + "/api/v1/environment-document", + ) + # Then + assert response.status_code == 401 + + +def test_get_environment_document_wrong_key( + client: TestClient, +) -> None: + # When + response = client.get( + "/api/v1/environment-document", headers={"X-Environment-Key": "ser.bad"} + ) + # Then + assert response.status_code == 401