Refactor

Author: matthewelwell

src/edge_proxy/environments.py

@@ -77,7 +77,7 @@ async def refresh_environment_caches(self):
     def get_flags_response_data(
         self, environment_key: str, feature: str = None
     ) -> dict[str, Any]:
-        environment_document = self.get_environment(client_side_key=environment_key)
+        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)
 
@@ -106,7 +106,7 @@ def get_flags_response_data(
     def get_identity_response_data(
         self, input_data: IdentityWithTraits, environment_key: str
     ) -> dict[str, Any]:
-        environment_document = self.get_environment(client_side_key=environment_key)
+        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)
 
@@ -140,21 +140,17 @@ def get_identity_response_data(
     def get_environment(
         self,
         *,
-        client_side_key: Optional[str] = None,
-        server_side_key: Optional[str] = None,
+        environment_key: Optional[str] = None,
     ) -> dict[str, Any]:
-        if client_side_key:
-            if environment_document := self.cache.get_environment(client_side_key):
-                return environment_document
-            raise FlagsmithUnknownKeyError(client_side_key)
-        if server_side_key:
-            for key_pair in self.settings.environment_key_pairs:
-                if key_pair.server_side_key == server_side_key:
-                    if environment_document := self.cache.get_environment(
-                        key_pair.client_side_key
-                    ):
-                        return environment_document
-            raise FlagsmithUnknownKeyError(server_side_key)
+        if 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(environment_key)
 
     async def _fetch_document(self, key_pair: EnvironmentKeyPair) -> dict[str, Any]:
         headers = {
@@ -198,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

@@ -130,7 +130,7 @@ async def environment_document(
     x_environment_key: str = Header(None),
 ) -> ORJSONResponse:
     if environment_doc := environment_service.get_environment(
-        server_side_key=x_environment_key,
+        environment_key=x_environment_key,
     ):
         return ORJSONResponse(environment_doc)
     return ORJSONResponse(status_code=401, content=None)

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(
-            client_side_key=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(
-            client_side_key=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(
-        client_side_key=settings.environment_key_pairs[0].client_side_key
+        environment_key=settings.environment_key_pairs[0].client_side_key
     )
     environment_service.get_environment(
-        client_side_key=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(client_side_key="test_env_key_unknown")
+        environment_service.get_environment(environment_key="test_env_key_unknown")
 
 
 @pytest.mark.asyncio