fix: reuse DynamoEnvironmentWrapper in DynamoIdentityWrapper (#6813)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: gagantrivedi

api/environments/dynamodb/migrator.py

@@ -60,7 +60,9 @@ def migrate(self):  # type: ignore[no-untyped-def]
         api_keys = EnvironmentAPIKey.objects.filter(environment__project_id=project_id)
         api_key_wrapper.write_api_keys(api_keys)
 
-        identity_wrapper = DynamoIdentityWrapper()
+        identity_wrapper = DynamoIdentityWrapper(
+            environment_wrapper=environment_wrapper
+        )
         identities = (
             Identity.objects.filter(environment__project__id=project_id)
             .select_related("environment")

api/environments/dynamodb/wrappers/identity_wrapper.py

@@ -37,6 +37,13 @@
 
 
 class DynamoIdentityWrapper(BaseDynamoWrapper):
+    def __init__(
+        self,
+        environment_wrapper: DynamoEnvironmentWrapper | None = None,
+    ) -> None:
+        super().__init__()
+        self._environment_wrapper = environment_wrapper or DynamoEnvironmentWrapper()
+
     def get_table_name(self) -> str | None:  # type: ignore[override]
         return settings.IDENTITIES_TABLE_NAME_DYNAMO
 
@@ -188,9 +195,8 @@ def get_segment_ids(
             identity = identity_model or IdentityModel.model_validate(
                 self.get_item_from_uuid(identity_pk)
             )
-            environment_wrapper = DynamoEnvironmentWrapper()
             environment = EnvironmentModel.model_validate(
-                environment_wrapper.get_item(identity.environment_api_key)
+                self._environment_wrapper.get_item(identity.environment_api_key)
             )
             context = map_environment_identity_to_context(
                 environment=environment,

api/environments/tasks.py

@@ -50,7 +50,7 @@ def delete_environment_from_dynamo(api_key: str, environment_id: str):  # type:
     environment_wrapper.delete_environment(api_key)
 
     # Delete identities
-    identity_wrapper = DynamoIdentityWrapper()
+    identity_wrapper = DynamoIdentityWrapper(environment_wrapper=environment_wrapper)
     identity_wrapper.delete_all_identities(api_key)
 
     # Delete environment_v2 documents

api/tests/unit/environments/dynamodb/test_unit_migrator.py

@@ -40,7 +40,9 @@ def test_migrate_calls_internal_methods_with_correct_arguments(  # type: ignore[
     identity_migrator.migrate()  # type: ignore[no-untyped-call]
 
     # Then
-    mocked_identity_wrapper.assert_called_with()
+    mocked_identity_wrapper.assert_called_with(
+        environment_wrapper=mocked_environment_wrapper.return_value
+    )
 
     args, kwargs = mocked_identity_wrapper.return_value.write_identities.call_args
     assert kwargs == {}

api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py

@@ -312,18 +312,19 @@ def test_get_segment_ids_returns_correct_segment_ids(  # type: ignore[no-untyped
     Segment.objects.create(name="Non matching segment", project=project)
 
     identity_document = map_identity_to_identity_document(identity)
-    dynamo_identity_wrapper = DynamoIdentityWrapper()
-    mocked_get_item_from_uuid = mocker.patch.object(
-        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
-    )
     identity_uuid = identity_document["identity_uuid"]
-
     environment_document = map_environment_to_environment_document(environment)
+
     mocked_environment_wrapper = mocker.patch(
         "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
     )
     mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
+    dynamo_identity_wrapper = DynamoIdentityWrapper()
+    mocked_get_item_from_uuid = mocker.patch.object(
+        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
+    )
+
     # When
     segment_ids = dynamo_identity_wrapper.get_segment_ids(identity_uuid)  # type: ignore[arg-type]
 
@@ -386,18 +387,19 @@ def test_get_segment_ids_with_segment_feature_overrides(
     )
 
     identity_document = map_identity_to_identity_document(identity)
-    dynamo_identity_wrapper = DynamoIdentityWrapper()
-    mocker.patch.object(
-        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
-    )
     identity_uuid = identity_document["identity_uuid"]
-
     environment_document = map_environment_to_environment_document(environment)
+
     mocked_environment_wrapper = mocker.patch(
         "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
     )
     mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
+    dynamo_identity_wrapper = DynamoIdentityWrapper()
+    mocker.patch.object(
+        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
+    )
+
     # When
     segment_ids = dynamo_identity_wrapper.get_segment_ids(identity_uuid)  # type: ignore[arg-type]
 
@@ -427,18 +429,19 @@ def test_get_segment_ids_returns_segment_using_in_operator_for_integer_traits(
     )
 
     identity_document = map_identity_to_identity_document(identity)
-    dynamo_identity_wrapper = DynamoIdentityWrapper()
-    mocker.patch.object(
-        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
-    )
     identity_uuid = identity_document["identity_uuid"]
-
     environment_document = map_environment_to_environment_document(environment)
+
     mocked_environment_wrapper = mocker.patch(
         "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
     )
     mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
+    dynamo_identity_wrapper = DynamoIdentityWrapper()
+    mocker.patch.object(
+        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
+    )
+
     # When
     segment_ids = dynamo_identity_wrapper.get_segment_ids(identity_uuid)  # type: ignore[arg-type]
 
@@ -490,20 +493,20 @@ def test_get_segment_ids_throws_value_error_if_arguments_not_valid():  # type: i
 
 def test_get_segment_ids_with_identity_model(identity, environment, mocker):  # type: ignore[no-untyped-def]
     # Given
-    dynamo_identity_wrapper = DynamoIdentityWrapper()
     identity_document = map_identity_to_identity_document(identity)
     identity_model = IdentityModel.parse_obj(identity_document)
-
-    mocker.patch.object(
-        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
-    )
-
     environment_document = map_environment_to_environment_document(environment)
+
     mocked_environment_wrapper = mocker.patch(
         "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
     )
     mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
+    dynamo_identity_wrapper = DynamoIdentityWrapper()
+    mocker.patch.object(
+        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
+    )
+
     # When
     segment_ids = dynamo_identity_wrapper.get_segment_ids(identity_model=identity_model)
 
@@ -640,6 +643,41 @@ def test_identity_wrapper__iter_all_items_paginated__capacity_budget_set__raises
     )
 
 
+def test_get_segment_ids__called_multiple_times__reuses_environment_wrapper(
+    project: "Project",
+    environment: "Environment",
+    identity: "Identity",
+    mocker: "MockerFixture",
+) -> None:
+    # Given
+    identity_document = map_identity_to_identity_document(identity)
+    environment_document = map_environment_to_environment_document(environment)
+
+    mocked_environment_wrapper_class = mocker.patch(
+        "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
+    )
+    mocked_environment_wrapper_class.return_value.get_item.return_value = (
+        environment_document
+    )
+
+    dynamo_identity_wrapper = DynamoIdentityWrapper()
+    mocker.patch.object(
+        dynamo_identity_wrapper,
+        "get_item_from_uuid",
+        return_value=identity_document,
+    )
+    identity_uuid = str(identity_document["identity_uuid"])
+
+    # When
+    dynamo_identity_wrapper.get_segment_ids(identity_uuid)
+    dynamo_identity_wrapper.get_segment_ids(identity_uuid)
+    dynamo_identity_wrapper.get_segment_ids(identity_uuid)
+
+    # Then - DynamoEnvironmentWrapper should be instantiated once
+    # (during DynamoIdentityWrapper.__init__), not once per get_segment_ids call.
+    assert mocked_environment_wrapper_class.call_count == 1
+
+
 def test_delete_all_identities__deletes_all_identities_documents_from_dynamodb(
     flagsmith_identities_table: Table,
     dynamodb_identity_wrapper: DynamoIdentityWrapper,