fix: Identity overrides are not deleted when deleting Edge identities (…

…#4460)

api/edge_api/identities/models.py

@@ -167,39 +167,64 @@ def remove_feature_override(self, feature_state: FeatureStateModel) -> None:
 
     def save(self, user: FFAdminUser = None, master_api_key: MasterAPIKey = None):
         self.dynamo_wrapper.put_item(self.to_document())
-        changes = self._get_changes()
-        if changes["feature_overrides"]:
+        changeset = self._get_changes()
+        self._update_feature_overrides(
+            changeset=changeset,
+            user=user,
+            master_api_key=master_api_key,
+        )
+        self._reset_initial_state()
+
+    def delete(
+        self, user: FFAdminUser = None, master_api_key: MasterAPIKey = None
+    ) -> None:
+        self.dynamo_wrapper.delete_item(self._engine_identity_model.composite_key)
+        self._engine_identity_model.identity_features.clear()
+        changeset = self._get_changes()
+        self._update_feature_overrides(
+            changeset=changeset,
+            user=user,
+            master_api_key=master_api_key,
+        )
+        self._reset_initial_state()
+
+    def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> None:
+        identity_feature_names = {
+            fs.feature.name for fs in self._engine_identity_model.identity_features
+        }
+        if not identity_feature_names.issubset(valid_feature_names):
+            self._engine_identity_model.prune_features(list(valid_feature_names))
+            sync_identity_document_features.delay(args=(str(self.identity_uuid),))
+
+    def to_document(self) -> dict:
+        return map_engine_identity_to_identity_document(self._engine_identity_model)
+
+    def _update_feature_overrides(
+        self,
+        changeset: IdentityChangeset,
+        user: FFAdminUser = None,
+        master_api_key: MasterAPIKey = None,
+    ) -> None:
+        if changeset["feature_overrides"]:
             # TODO: would this be simpler if we put a wrapper around FeatureStateModel instead?
             generate_audit_log_records.delay(
                 kwargs={
                     "environment_api_key": self.environment_api_key,
                     "identifier": self.identifier,
                     "user_id": getattr(user, "id", None),
-                    "changes": changes,
+                    "changes": changeset,
                     "identity_uuid": str(self.identity_uuid),
                     "master_api_key_id": getattr(master_api_key, "id", None),
                 }
             )
             update_flagsmith_environments_v2_identity_overrides.delay(
                 kwargs={
                     "environment_api_key": self.environment_api_key,
-                    "changes": changes,
+                    "changes": changeset,
                     "identity_uuid": str(self.identity_uuid),
                     "identifier": self.identifier,
                 }
             )
-        self._reset_initial_state()
-
-    def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> None:
-        identity_feature_names = {
-            fs.feature.name for fs in self._engine_identity_model.identity_features
-        }
-        if not identity_feature_names.issubset(valid_feature_names):
-            self._engine_identity_model.prune_features(list(valid_feature_names))
-            sync_identity_document_features.delay(args=(str(self.identity_uuid),))
-
-    def to_document(self) -> dict:
-        return map_engine_identity_to_identity_document(self._engine_identity_model)
 
     def _get_changes(self) -> IdentityChangeset:
         previous_instance = self._initial_state

api/edge_api/identities/views.py

@@ -160,7 +160,11 @@ def get_environment_from_request(self):
         )
 
     def perform_destroy(self, instance):
-        EdgeIdentity.dynamo_wrapper.delete_item(instance["composite_key"])
+        edge_identity = EdgeIdentity.from_identity_document(instance)
+        edge_identity.delete(
+            user=self.request.user,
+            master_api_key=getattr(self.request, "master_api_key", None),
+        )
 
     @swagger_auto_schema(
         responses={200: EdgeIdentityTraitsSerializer(many=True)},
@@ -281,7 +285,7 @@ def list(self, request, *args, **kwargs):
     def perform_destroy(self, instance):
         self.identity.remove_feature_override(instance)
         self.identity.save(
-            user=self.request.user.id,
+            user=self.request.user,
             master_api_key=getattr(self.request, "master_api_key", None),
         )
 

api/tests/integration/edge_api/identities/conftest.py

@@ -1,8 +1,40 @@
+from typing import Any
+
 import pytest
+from boto3.dynamodb.conditions import Key
+from flag_engine.identities.models import IdentityModel
+
+from edge_api.identities.models import EdgeIdentity
+from environments.dynamodb.wrappers.environment_wrapper import (
+    DynamoEnvironmentV2Wrapper,
+)
 
 
 @pytest.fixture()
 def webhook_mock(mocker):
     return mocker.patch(
         "edge_api.identities.serializers.call_environment_webhook_for_feature_state_change"
     )
+
+
+@pytest.fixture()
+def identity_overrides_v2(
+    dynamo_enabled_environment: int,
+    identity_document_without_fs: dict[str, Any],
+    identity_document: dict[str, Any],
+    dynamodb_wrapper_v2: DynamoEnvironmentV2Wrapper,
+) -> list[str]:
+    edge_identity = EdgeIdentity.from_identity_document(identity_document_without_fs)
+    for feature_override in IdentityModel.model_validate(
+        identity_document
+    ).identity_features:
+        edge_identity.add_feature_override(feature_override)
+    edge_identity.save()
+    return [
+        item["document_key"]
+        for item in dynamodb_wrapper_v2.query_get_all_items(
+            KeyConditionExpression=Key("environment_id").eq(
+                str(dynamo_enabled_environment)
+            ),
+        )
+    ]

api/tests/integration/edge_api/identities/test_edge_identity_viewset.py

@@ -1,11 +1,20 @@
 import json
 import urllib
+from typing import Any
 
+from boto3.dynamodb.conditions import Key
 from django.urls import reverse
 from rest_framework import status
 from rest_framework.exceptions import NotFound
+from rest_framework.test import APIClient
 
 from edge_api.identities.views import EdgeIdentityViewSet
+from environments.dynamodb.wrappers.environment_wrapper import (
+    DynamoEnvironmentV2Wrapper,
+)
+from environments.dynamodb.wrappers.identity_wrapper import (
+    DynamoIdentityWrapper,
+)
 
 
 def test_get_identities_returns_bad_request_if_dynamo_is_not_enabled(
@@ -125,33 +134,38 @@ def test_create_identity_returns_400_if_identity_already_exists(
 
 
 def test_delete_identity(
-    admin_client,
-    dynamo_enabled_environment,
-    environment_api_key,
-    identity_document,
-    edge_identity_dynamo_wrapper_mock,
-):
+    admin_client: APIClient,
+    dynamo_enabled_environment: int,
+    environment_api_key: str,
+    identity_document_without_fs: dict[str, Any],
+    identity_document: dict[str, Any],
+    dynamodb_identity_wrapper: DynamoIdentityWrapper,
+    dynamodb_wrapper_v2: DynamoEnvironmentV2Wrapper,
+    identity_overrides_v2: list[str],
+) -> None:
     # Given
     identity_uuid = identity_document["identity_uuid"]
     url = reverse(
         "api-v1:environments:environment-edge-identities-detail",
         args=[environment_api_key, identity_uuid],
     )
 
-    edge_identity_dynamo_wrapper_mock.get_item_from_uuid_or_404.return_value = (
-        identity_document
-    )
     # When
     response = admin_client.delete(url)
 
     # Then
     assert response.status_code == status.HTTP_204_NO_CONTENT
 
-    edge_identity_dynamo_wrapper_mock.get_item_from_uuid_or_404.assert_called_with(
-        identity_uuid
-    )
-    edge_identity_dynamo_wrapper_mock.delete_item.assert_called_with(
-        identity_document["composite_key"]
+    assert not dynamodb_identity_wrapper.query_items(
+        IndexName="identity_uuid-index",
+        KeyConditionExpression=Key("identity_uuid").eq(identity_uuid),
+    )["Count"]
+    assert not list(
+        dynamodb_wrapper_v2.query_get_all_items(
+            KeyConditionExpression=Key("environment_id").eq(
+                str(dynamo_enabled_environment)
+            )
+        )
     )