feat: Add environment document endpoint (#150)

* add environment-document endpoint

* nit

* remove prints

* separate environment document test cases

* Abstract environment-document into service, make args explicit

* lint

* Refactor

* Refactor

---------

Co-authored-by: Matthew Elwell <[email protected]>

Author: rolodato

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)

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,

tests/test_environments.py

@@ -111,32 +111,32 @@ 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
     )
     assert mock_client.get.call_count == 2
 
     # 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
 
 
 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()

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