From b7ecd75810d0d98221b775a7f87ba1e73b98647a Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Mon, 8 Jan 2024 17:40:22 +0530 Subject: [PATCH] feat(dynamo_documents): propagate delete to dynamo (#3220) --- api/conftest.py | 23 --- .../identities/edge_identity_service.py | 2 +- api/edge_api/identities/tasks.py | 2 +- api/environments/dynamodb/__init__.py | 4 +- api/environments/dynamodb/migrator.py | 4 +- api/environments/dynamodb/services.py | 7 +- api/environments/dynamodb/types.py | 4 + .../dynamodb/wrappers/__init__.py | 13 ++ api/environments/dynamodb/wrappers/base.py | 45 +++++ .../wrappers/environment_api_key_wrapper.py | 31 ++++ .../dynamodb/wrappers/environment_wrapper.py | 129 +++++++++++++ .../identity_wrapper.py} | 175 +++--------------- api/environments/models.py | 23 ++- api/environments/tasks.py | 24 ++- api/projects/models.py | 6 + api/tests/conftest.py | 89 +++++++++ .../dynamodb/test_unit_services.py | 2 +- .../test_unit_dynamodb_project_metadata.py | 28 +++ .../dynamodb/wrappers/__init__.py | 0 .../test_unit_dynamo_environment_wrapper.py | 17 ++ ...it_dynamodb_environment_api_key_wrapper.py | 0 ...st_unit_dynamodb_environment_v2_wrapper.py | 44 ++++- .../test_unit_dynamodb_identity_wrapper.py | 65 +++++-- .../test_unit_environments_models.py | 78 +++++++- .../test_unit_environments_tasks.py | 33 ++++ .../test_unit_features_features_service.py | 2 +- 26 files changed, 642 insertions(+), 208 deletions(-) create mode 100644 api/environments/dynamodb/wrappers/__init__.py create mode 100644 api/environments/dynamodb/wrappers/base.py create mode 100644 api/environments/dynamodb/wrappers/environment_api_key_wrapper.py create mode 100644 api/environments/dynamodb/wrappers/environment_wrapper.py rename api/environments/dynamodb/{dynamodb_wrapper.py => wrappers/identity_wrapper.py} (50%) create mode 100644 api/tests/unit/environments/dynamodb/wrappers/__init__.py rename api/tests/unit/environments/dynamodb/{ => wrappers}/test_unit_dynamo_environment_wrapper.py (78%) rename api/tests/unit/environments/dynamodb/{ => wrappers}/test_unit_dynamodb_environment_api_key_wrapper.py (100%) rename api/tests/unit/environments/dynamodb/{ => wrappers}/test_unit_dynamodb_environment_v2_wrapper.py (78%) rename api/tests/unit/environments/dynamodb/{ => wrappers}/test_unit_dynamodb_identity_wrapper.py (86%) diff --git a/api/conftest.py b/api/conftest.py index dd3c25ea94de..c4db0c583ccb 100644 --- a/api/conftest.py +++ b/api/conftest.py @@ -8,15 +8,10 @@ from flag_engine.segments.constants import EQUAL from moto import mock_dynamodb from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table -from pytest_django.fixtures import SettingsWrapper from rest_framework.authtoken.models import Token from rest_framework.test import APIClient from api_keys.models import MasterAPIKey -from environments.dynamodb.dynamodb_wrapper import ( - DynamoEnvironmentV2Wrapper, - DynamoIdentityWrapper, -) from environments.identities.models import Identity from environments.identities.traits.models import Trait from environments.models import Environment, EnvironmentAPIKey @@ -618,21 +613,3 @@ def flagsmith_environments_v2_table(dynamodb: DynamoDBServiceResource) -> Table: ], BillingMode="PAY_PER_REQUEST", ) - - -@pytest.fixture -def dynamodb_identity_wrapper( - settings: SettingsWrapper, - flagsmith_identities_table: Table, -) -> DynamoIdentityWrapper: - settings.IDENTITIES_TABLE_NAME_DYNAMO = flagsmith_identities_table.name - return DynamoIdentityWrapper() - - -@pytest.fixture -def dynamodb_wrapper_v2( - settings: SettingsWrapper, - flagsmith_environments_v2_table: Table, -) -> DynamoEnvironmentV2Wrapper: - settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = flagsmith_environments_v2_table.name - return DynamoEnvironmentV2Wrapper() diff --git a/api/edge_api/identities/edge_identity_service.py b/api/edge_api/identities/edge_identity_service.py index 010ffd204cdc..ef2b471f7c73 100644 --- a/api/edge_api/identities/edge_identity_service.py +++ b/api/edge_api/identities/edge_identity_service.py @@ -1,6 +1,6 @@ import typing -from environments.dynamodb.dynamodb_wrapper import DynamoEnvironmentV2Wrapper +from environments.dynamodb import DynamoEnvironmentV2Wrapper from environments.dynamodb.types import IdentityOverrideV2 ddb_environment_v2_wrapper = DynamoEnvironmentV2Wrapper() diff --git a/api/edge_api/identities/tasks.py b/api/edge_api/identities/tasks.py index fb99729c9c54..48d21dcb7829 100644 --- a/api/edge_api/identities/tasks.py +++ b/api/edge_api/identities/tasks.py @@ -6,7 +6,7 @@ from audit.models import AuditLog from audit.related_object_type import RelatedObjectType from edge_api.identities.types import IdentityChangeset -from environments.dynamodb.dynamodb_wrapper import DynamoEnvironmentV2Wrapper +from environments.dynamodb import DynamoEnvironmentV2Wrapper from environments.models import Environment, Webhook from features.models import Feature, FeatureState from task_processor.decorators import register_task_handler diff --git a/api/environments/dynamodb/__init__.py b/api/environments/dynamodb/__init__.py index f36c2a219915..26348bce0552 100644 --- a/api/environments/dynamodb/__init__.py +++ b/api/environments/dynamodb/__init__.py @@ -1,4 +1,5 @@ -from .dynamodb_wrapper import ( +from .types import DynamoProjectMetadata +from .wrappers import ( DynamoEnvironmentAPIKeyWrapper, DynamoEnvironmentV2Wrapper, DynamoEnvironmentWrapper, @@ -10,4 +11,5 @@ "DynamoEnvironmentV2Wrapper", "DynamoEnvironmentWrapper", "DynamoIdentityWrapper", + "DynamoProjectMetadata", ) diff --git a/api/environments/dynamodb/migrator.py b/api/environments/dynamodb/migrator.py index 7367a11190ce..d53763c20160 100644 --- a/api/environments/dynamodb/migrator.py +++ b/api/environments/dynamodb/migrator.py @@ -8,12 +8,12 @@ from projects.models import Project from util.queryset import iterator_with_prefetch -from .dynamodb_wrapper import ( +from .types import DynamoProjectMetadata, ProjectIdentityMigrationStatus +from .wrappers import ( DynamoEnvironmentAPIKeyWrapper, DynamoEnvironmentWrapper, DynamoIdentityWrapper, ) -from .types import DynamoProjectMetadata, ProjectIdentityMigrationStatus class IdentityMigrator: diff --git a/api/environments/dynamodb/services.py b/api/environments/dynamodb/services.py index 10af48ce855c..23c9d9ecd48e 100644 --- a/api/environments/dynamodb/services.py +++ b/api/environments/dynamodb/services.py @@ -1,7 +1,9 @@ import logging from typing import Generator, Iterable -from environments.dynamodb.dynamodb_wrapper import ( +from flag_engine.identities.models import IdentityModel + +from environments.dynamodb import ( DynamoEnvironmentV2Wrapper, DynamoIdentityWrapper, ) @@ -53,9 +55,10 @@ def _iter_paginated_overrides( ) -> Generator[IdentityOverrideV2, None, None]: for environment in environments: environment_api_key = environment.api_key - for identity in identity_wrapper.iter_all_items_paginated( + for item in identity_wrapper.iter_all_items_paginated( environment_api_key=environment_api_key, ): + identity = IdentityModel.parse_obj(item) for feature_state in identity.identity_features: yield map_engine_feature_state_to_identity_override( feature_state=feature_state, diff --git a/api/environments/dynamodb/types.py b/api/environments/dynamodb/types.py index 88b6c192a6aa..552321777fec 100644 --- a/api/environments/dynamodb/types.py +++ b/api/environments/dynamodb/types.py @@ -72,6 +72,10 @@ def finish_identity_migration(self): def _save(self): return project_metadata_table.put_item(Item=asdict(self)) + def delete(self): + if project_metadata_table: + project_metadata_table.delete_item(Key={"id": self.id}) + class IdentityOverrideV2(BaseModel): environment_id: str diff --git a/api/environments/dynamodb/wrappers/__init__.py b/api/environments/dynamodb/wrappers/__init__.py new file mode 100644 index 000000000000..b5b810a890de --- /dev/null +++ b/api/environments/dynamodb/wrappers/__init__.py @@ -0,0 +1,13 @@ +from .environment_api_key_wrapper import DynamoEnvironmentAPIKeyWrapper +from .environment_wrapper import ( + DynamoEnvironmentV2Wrapper, + DynamoEnvironmentWrapper, +) +from .identity_wrapper import DynamoIdentityWrapper + +__all__ = ( + "DynamoEnvironmentAPIKeyWrapper", + "DynamoEnvironmentV2Wrapper", + "DynamoEnvironmentWrapper", + "DynamoIdentityWrapper", +) diff --git a/api/environments/dynamodb/wrappers/base.py b/api/environments/dynamodb/wrappers/base.py new file mode 100644 index 000000000000..abac7ca9b442 --- /dev/null +++ b/api/environments/dynamodb/wrappers/base.py @@ -0,0 +1,45 @@ +import typing + +import boto3 +from botocore.config import Config + +if typing.TYPE_CHECKING: + from mypy_boto3_dynamodb.service_resource import Table + + +class BaseDynamoWrapper: + table_name: str = None + + def __init__(self) -> None: + self._table: typing.Optional["Table"] = None + + @property + def table(self) -> typing.Optional["Table"]: + if not self._table: + self._table = self.get_table() + return self._table + + def get_table_name(self) -> str: + return self.table_name + + def get_table(self) -> typing.Optional["Table"]: + if table_name := self.get_table_name(): + return boto3.resource("dynamodb", config=Config(tcp_keepalive=True)).Table( + table_name + ) + + @property + def is_enabled(self) -> bool: + return self.table is not None + + def query_get_all_items(self, **kwargs: dict) -> typing.Generator[dict, None, None]: + while True: + query_response = self.table.query(**kwargs) + for item in query_response["Items"]: + yield item + + last_evaluated_key = query_response.get("LastEvaluatedKey") + if not last_evaluated_key: + break + + kwargs["ExclusiveStartKey"] = last_evaluated_key diff --git a/api/environments/dynamodb/wrappers/environment_api_key_wrapper.py b/api/environments/dynamodb/wrappers/environment_api_key_wrapper.py new file mode 100644 index 000000000000..c95ebd787d2b --- /dev/null +++ b/api/environments/dynamodb/wrappers/environment_api_key_wrapper.py @@ -0,0 +1,31 @@ +import typing + +from django.conf import settings + +from util.mappers import ( + map_environment_api_key_to_environment_api_key_document, +) + +from .base import BaseDynamoWrapper + +if typing.TYPE_CHECKING: + from environments.models import EnvironmentAPIKey + + +class DynamoEnvironmentAPIKeyWrapper(BaseDynamoWrapper): + table_name = settings.ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO + + def write_api_key(self, api_key: "EnvironmentAPIKey"): + self.write_api_keys([api_key]) + + def write_api_keys(self, api_keys: typing.Iterable["EnvironmentAPIKey"]): + with self.table.batch_writer() as writer: + for api_key in api_keys: + writer.put_item( + Item=map_environment_api_key_to_environment_api_key_document( + api_key + ) + ) + + def delete_api_key(self, key: str) -> None: + self.table.delete_item(Key={"key": key}) diff --git a/api/environments/dynamodb/wrappers/environment_wrapper.py b/api/environments/dynamodb/wrappers/environment_wrapper.py new file mode 100644 index 000000000000..c8f40d2536d6 --- /dev/null +++ b/api/environments/dynamodb/wrappers/environment_wrapper.py @@ -0,0 +1,129 @@ +import typing +from typing import Any, Iterable + +from boto3.dynamodb.conditions import Key +from django.conf import settings +from django.core.exceptions import ObjectDoesNotExist + +from environments.dynamodb.constants import ( + DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT, + ENVIRONMENTS_V2_PARTITION_KEY, + ENVIRONMENTS_V2_SORT_KEY, +) +from environments.dynamodb.types import IdentityOverridesV2Changeset +from environments.dynamodb.utils import ( + get_environments_v2_identity_override_document_key, +) +from util.mappers import ( + map_environment_to_environment_document, + map_environment_to_environment_v2_document, + map_identity_override_to_identity_override_document, +) +from util.util import iter_paired_chunks + +from .base import BaseDynamoWrapper + +if typing.TYPE_CHECKING: + from mypy_boto3_dynamodb.type_defs import QueryInputRequestTypeDef + + from environments.models import Environment + + +class BaseDynamoEnvironmentWrapper(BaseDynamoWrapper): + def write_environment(self, environment: "Environment") -> None: + self.write_environments([environment]) + + def write_environments(self, environments: Iterable["Environment"]) -> None: + raise NotImplementedError() + + +class DynamoEnvironmentWrapper(BaseDynamoEnvironmentWrapper): + table_name = settings.ENVIRONMENTS_TABLE_NAME_DYNAMO + + def write_environments(self, environments: Iterable["Environment"]): + with self.table.batch_writer() as writer: + for environment in environments: + writer.put_item( + Item=map_environment_to_environment_document(environment), + ) + + def get_item(self, api_key: str) -> dict: + try: + return self.table.get_item(Key={"api_key": api_key})["Item"] + except KeyError as e: + raise ObjectDoesNotExist() from e + + def delete_environment(self, api_key: str) -> None: + self.table.delete_item(Key={"api_key": api_key}) + + +class DynamoEnvironmentV2Wrapper(BaseDynamoEnvironmentWrapper): + def get_table_name(self) -> str | None: + return settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO + + def get_identity_overrides_by_environment_id( + self, + environment_id: int, + feature_id: int | None = None, + ) -> typing.List[dict[str, Any]]: + try: + response = self.table.query( + KeyConditionExpression=Key(ENVIRONMENTS_V2_PARTITION_KEY).eq( + str(environment_id), + ) + & Key(ENVIRONMENTS_V2_SORT_KEY).begins_with( + get_environments_v2_identity_override_document_key( + feature_id=feature_id, + ), + ) + ) + return response["Items"] + except KeyError as e: + raise ObjectDoesNotExist() from e + + def update_identity_overrides( + self, + changeset: IdentityOverridesV2Changeset, + ) -> None: + for to_put, to_delete in iter_paired_chunks( + changeset.to_put, + changeset.to_delete, + chunk_size=DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT, + ): + with self.table.batch_writer() as writer: + for identity_override_to_delete in to_delete: + writer.delete_item( + Key={ + ENVIRONMENTS_V2_PARTITION_KEY: identity_override_to_delete.environment_id, + ENVIRONMENTS_V2_SORT_KEY: identity_override_to_delete.document_key, + }, + ) + for identity_override_to_put in to_put: + writer.put_item( + Item=map_identity_override_to_identity_override_document( + identity_override_to_put + ), + ) + + def write_environments(self, environments: Iterable["Environment"]) -> None: + with self.table.batch_writer() as writer: + for environment in environments: + writer.put_item( + Item=map_environment_to_environment_v2_document(environment), + ) + + def delete_environment(self, environment_id: int): + environment_id = str(environment_id) + filter_expression = Key(ENVIRONMENTS_V2_PARTITION_KEY).eq(environment_id) + query_kwargs: "QueryInputRequestTypeDef" = { + "KeyConditionExpression": filter_expression, + "ProjectionExpression": "document_key", + } + with self.table.batch_writer() as writer: + for item in self.query_get_all_items(**query_kwargs): + writer.delete_item( + Key={ + ENVIRONMENTS_V2_PARTITION_KEY: environment_id, + ENVIRONMENTS_V2_SORT_KEY: item["document_key"], + }, + ) diff --git a/api/environments/dynamodb/dynamodb_wrapper.py b/api/environments/dynamodb/wrappers/identity_wrapper.py similarity index 50% rename from api/environments/dynamodb/dynamodb_wrapper.py rename to api/environments/dynamodb/wrappers/identity_wrapper.py index 208c2bed7728..6351419d0fc7 100644 --- a/api/environments/dynamodb/dynamodb_wrapper.py +++ b/api/environments/dynamodb/wrappers/identity_wrapper.py @@ -1,11 +1,9 @@ import logging import typing from contextlib import suppress -from typing import Any, Iterable +from typing import Iterable -import boto3 from boto3.dynamodb.conditions import Key -from botocore.config import Config from django.conf import settings from django.core.exceptions import ObjectDoesNotExist from flag_engine.environments.builders import build_environment_model @@ -14,66 +12,24 @@ from flag_engine.segments.evaluator import get_identity_segments from rest_framework.exceptions import NotFound +from environments.dynamodb.constants import IDENTITIES_PAGINATION_LIMIT +from util.mappers import map_identity_to_identity_document + +from .base import BaseDynamoWrapper +from .environment_wrapper import DynamoEnvironmentWrapper + if typing.TYPE_CHECKING: - from mypy_boto3_dynamodb.service_resource import Table from mypy_boto3_dynamodb.type_defs import ( + QueryInputRequestTypeDef, QueryOutputTableTypeDef, TableAttributeValueTypeDef, - QueryInputRequestTypeDef, ) -from environments.dynamodb.constants import ( - DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT, - ENVIRONMENTS_V2_PARTITION_KEY, - ENVIRONMENTS_V2_SORT_KEY, - IDENTITIES_PAGINATION_LIMIT, -) -from environments.dynamodb.types import IdentityOverridesV2Changeset -from environments.dynamodb.utils import ( - get_environments_v2_identity_override_document_key, -) -from util.mappers import ( - map_environment_api_key_to_environment_api_key_document, - map_environment_to_environment_document, - map_environment_to_environment_v2_document, - map_identity_override_to_identity_override_document, - map_identity_to_identity_document, -) -from util.util import iter_paired_chunks - -if typing.TYPE_CHECKING: from environments.identities.models import Identity - from environments.models import Environment, EnvironmentAPIKey logger = logging.getLogger() -class BaseDynamoWrapper: - table_name: str = None - - def __init__(self) -> None: - self._table: typing.Optional["Table"] = None - - @property - def table(self) -> typing.Optional["Table"]: - if not self._table: - self._table = self.get_table() - return self._table - - def get_table_name(self) -> str: - return self.table_name - - def get_table(self) -> typing.Optional["Table"]: - if table_name := self.get_table_name(): - return boto3.resource("dynamodb", config=Config(tcp_keepalive=True)).Table( - table_name - ) - - @property - def is_enabled(self) -> bool: - return self.table is not None - - class DynamoIdentityWrapper(BaseDynamoWrapper): def get_table_name(self) -> str | None: return settings.IDENTITIES_TABLE_NAME_DYNAMO @@ -103,6 +59,14 @@ def get_item(self, composite_key: str) -> typing.Optional[dict]: def delete_item(self, composite_key: str): self.table.delete_item(Key={"composite_key": composite_key}) + def delete_all_identities(self, environment_api_key: str): + with self.table.batch_writer() as writer: + for item in self.iter_all_items_paginated( + environment_api_key=environment_api_key, + projection_expression="composite_key", + ): + writer.delete_item(Key={"composite_key": item["composite_key"]}) + def get_item_from_uuid(self, uuid: str) -> dict: filter_expression = Key("identity_uuid").eq(uuid) query_kwargs = { @@ -126,6 +90,7 @@ def get_all_items( environment_api_key: str, limit: int, start_key: dict[str, "TableAttributeValueTypeDef"] | None = None, + projection_expression: str = None, ) -> "QueryOutputTableTypeDef": filter_expression = Key("environment_api_key").eq(environment_api_key) query_kwargs: "QueryInputRequestTypeDef" = { @@ -133,6 +98,9 @@ def get_all_items( "KeyConditionExpression": filter_expression, "Limit": limit, } + if projection_expression: + query_kwargs["ProjectionExpression"] = projection_expression + if start_key: query_kwargs["ExclusiveStartKey"] = start_key return self.query_items(**query_kwargs) @@ -141,18 +109,20 @@ def iter_all_items_paginated( self, environment_api_key: str, limit: int = IDENTITIES_PAGINATION_LIMIT, - ) -> typing.Generator[IdentityModel, None, None]: + projection_expression: str = None, + ) -> typing.Generator[dict, None, None]: last_evaluated_key = "initial" get_all_items_kwargs = { "environment_api_key": environment_api_key, "limit": limit, + "projection_expression": projection_expression, } while last_evaluated_key: query_response = self.get_all_items( **get_all_items_kwargs, ) for item in query_response["Items"]: - yield IdentityModel.parse_obj(item) + yield item if last_evaluated_key := query_response.get("LastEvaluatedKey"): get_all_items_kwargs["start_key"] = last_evaluated_key @@ -194,100 +164,3 @@ def get_segment_ids( return [segment.id for segment in segments] return [] - - -class BaseDynamoEnvironmentWrapper(BaseDynamoWrapper): - def write_environment(self, environment: "Environment") -> None: - self.write_environments([environment]) - - def write_environments(self, environments: Iterable["Environment"]) -> None: - raise NotImplementedError() - - -class DynamoEnvironmentWrapper(BaseDynamoEnvironmentWrapper): - table_name = settings.ENVIRONMENTS_TABLE_NAME_DYNAMO - - def write_environments(self, environments: Iterable["Environment"]): - with self.table.batch_writer() as writer: - for environment in environments: - writer.put_item( - Item=map_environment_to_environment_document(environment), - ) - - def get_item(self, api_key: str) -> dict: - try: - return self.table.get_item(Key={"api_key": api_key})["Item"] - except KeyError as e: - raise ObjectDoesNotExist() from e - - -class DynamoEnvironmentV2Wrapper(BaseDynamoEnvironmentWrapper): - def get_table_name(self) -> str | None: - return settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO - - def get_identity_overrides_by_environment_id( - self, - environment_id: int, - feature_id: int | None = None, - ) -> typing.List[dict[str, Any]]: - try: - response = self.table.query( - KeyConditionExpression=Key(ENVIRONMENTS_V2_PARTITION_KEY).eq( - str(environment_id), - ) - & Key(ENVIRONMENTS_V2_SORT_KEY).begins_with( - get_environments_v2_identity_override_document_key( - feature_id=feature_id, - ), - ) - ) - return response["Items"] - except KeyError as e: - raise ObjectDoesNotExist() from e - - def update_identity_overrides( - self, - changeset: IdentityOverridesV2Changeset, - ) -> None: - for to_put, to_delete in iter_paired_chunks( - changeset.to_put, - changeset.to_delete, - chunk_size=DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT, - ): - with self.table.batch_writer() as writer: - for identity_override_to_delete in to_delete: - writer.delete_item( - Key={ - ENVIRONMENTS_V2_PARTITION_KEY: identity_override_to_delete.environment_id, - ENVIRONMENTS_V2_SORT_KEY: identity_override_to_delete.document_key, - }, - ) - for identity_override_to_put in to_put: - writer.put_item( - Item=map_identity_override_to_identity_override_document( - identity_override_to_put - ), - ) - - def write_environments(self, environments: Iterable["Environment"]) -> None: - with self.table.batch_writer() as writer: - for environment in environments: - writer.put_item( - Item=map_environment_to_environment_v2_document(environment), - ) - - -class DynamoEnvironmentAPIKeyWrapper(BaseDynamoWrapper): - table_name = settings.ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO - - def write_api_key(self, api_key: "EnvironmentAPIKey"): - self.write_api_keys([api_key]) - - def write_api_keys(self, api_keys: Iterable["EnvironmentAPIKey"]): - with self.table.batch_writer() as writer: - for api_key in api_keys: - writer.put_item( - Item=map_environment_api_key_to_environment_api_key_document( - api_key - ) - ) diff --git a/api/environments/models.py b/api/environments/models.py index 657810aaa9cd..6497aa3e4155 100644 --- a/api/environments/models.py +++ b/api/environments/models.py @@ -16,6 +16,7 @@ from django.utils.translation import ugettext_lazy as _ from django_lifecycle import ( AFTER_CREATE, + AFTER_DELETE, AFTER_SAVE, AFTER_UPDATE, BEFORE_UPDATE, @@ -147,6 +148,13 @@ def clear_environment_cache(self): def validate_use_v2_feature_versioning(self): raise FeatureVersioningError("Cannot revert from v2 feature versioning.") + @hook(AFTER_DELETE) + def delete_from_dynamo(self): + if self.project.enable_dynamo_db and environment_wrapper.is_enabled: + from environments.tasks import delete_environment_from_dynamo + + delete_environment_from_dynamo.delay(args=(self.api_key, self.id)) + def __str__(self): return "Project %s - Environment %s" % (self.project.name, self.name) @@ -449,10 +457,17 @@ def natural_key(self): def is_valid(self) -> bool: return self.active and (not self.expires_at or self.expires_at > timezone.now()) - @hook(AFTER_SAVE) + @hook(AFTER_SAVE, when="_should_update_dynamo", is_now=True) def send_to_dynamo(self): - if ( + environment_api_key_wrapper.write_api_key(self) + + @hook(AFTER_DELETE, when="_should_update_dynamo", is_now=True) + def delete_from_dynamo(self): + environment_api_key_wrapper.delete_api_key(self.key) + + @property + def _should_update_dynamo(self) -> bool: + return ( self.environment.project.enable_dynamo_db and environment_api_key_wrapper.is_enabled - ): - environment_api_key_wrapper.write_api_key(self) + ) diff --git a/api/environments/tasks.py b/api/environments/tasks.py index 6197b74ed3db..e148834d5382 100644 --- a/api/environments/tasks.py +++ b/api/environments/tasks.py @@ -1,6 +1,13 @@ from audit.models import AuditLog -from environments.dynamodb import DynamoEnvironmentWrapper -from environments.models import Environment +from environments.dynamodb import ( + DynamoEnvironmentWrapper, + DynamoIdentityWrapper, +) +from environments.models import ( + Environment, + environment_v2_wrapper, + environment_wrapper, +) from sse import ( send_environment_update_message_for_environment, send_environment_update_message_for_project, @@ -31,3 +38,16 @@ def process_environment_update(audit_log_id: int): send_environment_update_message_for_environment(audit_log.environment) else: send_environment_update_message_for_project(audit_log.project) + + +@register_task_handler() +def delete_environment_from_dynamo(api_key: str, environment_id: str): + # Delete environment + environment_wrapper.delete_environment(api_key) + + # Delete identities + identity_wrapper = DynamoIdentityWrapper() + identity_wrapper.delete_all_identities(api_key) + + # Delete environment_v2 documents + environment_v2_wrapper.delete_environment(environment_id) diff --git a/api/projects/models.py b/api/projects/models.py index 639bab563eb9..16c137c057bb 100644 --- a/api/projects/models.py +++ b/api/projects/models.py @@ -10,6 +10,7 @@ from django.db.models import Count from django.utils import timezone from django_lifecycle import ( + AFTER_DELETE, AFTER_SAVE, AFTER_UPDATE, BEFORE_CREATE, @@ -17,6 +18,7 @@ hook, ) +from environments.dynamodb import DynamoProjectMetadata from organisations.models import Organisation from permissions.models import ( PROJECT_PERMISSION_TYPE, @@ -152,6 +154,10 @@ def trigger_environments_v2_migration(self) -> None: def write_to_dynamo(self): write_environments_to_dynamodb.delay(kwargs={"project_id": self.id}) + @hook(AFTER_DELETE) + def clean_up_dynamo(self): + DynamoProjectMetadata(self.id).delete() + @property def is_edge_project_by_default(self) -> bool: return bool( diff --git a/api/tests/conftest.py b/api/tests/conftest.py index 365fdaf16d66..ac7d67a7af0a 100644 --- a/api/tests/conftest.py +++ b/api/tests/conftest.py @@ -1,4 +1,14 @@ import pytest +from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table +from pytest_django.fixtures import SettingsWrapper + +from environments.dynamodb import ( + DynamoEnvironmentV2Wrapper, + DynamoEnvironmentWrapper, + DynamoIdentityWrapper, +) +from environments.models import Environment +from util.mappers import map_environment_to_environment_document @pytest.fixture() @@ -6,3 +16,82 @@ def edge_identity_dynamo_wrapper_mock(mocker): return mocker.patch( "edge_api.identities.models.EdgeIdentity.dynamo_wrapper", ) + + +@pytest.fixture() +def flagsmith_environment_api_key_table(dynamodb: "DynamoDBServiceResource") -> "Table": + return dynamodb.create_table( + TableName="flagsmith_environment_api_key", + KeySchema=[{"AttributeName": "key", "KeyType": "HASH"}], + AttributeDefinitions=[ + {"AttributeName": "key", "AttributeType": "S"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + + +@pytest.fixture() +def flagsmith_environment_table(dynamodb: "DynamoDBServiceResource") -> "Table": + return dynamodb.create_table( + TableName="flagsmith_environments", + KeySchema=[{"AttributeName": "api_key", "KeyType": "HASH"}], + AttributeDefinitions=[ + {"AttributeName": "api_key", "AttributeType": "S"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + + +@pytest.fixture() +def flagsmith_project_metadata_table(dynamodb: "DynamoDBServiceResource") -> "Table": + return dynamodb.create_table( + TableName="flagsmith_project_metadata", + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=[ + {"AttributeName": "id", "AttributeType": "N"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + + +@pytest.fixture +def dynamodb_identity_wrapper( + settings: SettingsWrapper, + flagsmith_identities_table: Table, +) -> DynamoIdentityWrapper: + settings.IDENTITIES_TABLE_NAME_DYNAMO = flagsmith_identities_table.name + return DynamoIdentityWrapper() + + +@pytest.fixture +def dynamodb_wrapper_v2( + settings: SettingsWrapper, + flagsmith_environments_v2_table: Table, +) -> DynamoEnvironmentV2Wrapper: + settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = flagsmith_environments_v2_table.name + return DynamoEnvironmentV2Wrapper() + + +@pytest.fixture() +def dynamo_enabled_project_environment_one_document( + flagsmith_environment_table: Table, + dynamo_enabled_project_environment_one: Environment, +) -> dict: + environment_dict = map_environment_to_environment_document( + dynamo_enabled_project_environment_one + ) + + flagsmith_environment_table.put_item( + Item=environment_dict, + ) + return environment_dict + + +@pytest.fixture() +def dynamo_environment_wrapper( + flagsmith_environment_table: Table, +) -> DynamoEnvironmentWrapper: + wrapper = DynamoEnvironmentWrapper() + + wrapper.table_name = flagsmith_environment_table.name + return wrapper diff --git a/api/tests/unit/environments/dynamodb/test_unit_services.py b/api/tests/unit/environments/dynamodb/test_unit_services.py index 63fdd71bb152..6314ec4a77b8 100644 --- a/api/tests/unit/environments/dynamodb/test_unit_services.py +++ b/api/tests/unit/environments/dynamodb/test_unit_services.py @@ -1,7 +1,7 @@ from mypy_boto3_dynamodb.service_resource import Table from pytest_mock import MockerFixture -from environments.dynamodb.dynamodb_wrapper import ( +from environments.dynamodb import ( DynamoEnvironmentV2Wrapper, DynamoIdentityWrapper, ) diff --git a/api/tests/unit/environments/dynamodb/types/test_unit_dynamodb_project_metadata.py b/api/tests/unit/environments/dynamodb/types/test_unit_dynamodb_project_metadata.py index a4f232081686..b3d9e8f709f9 100644 --- a/api/tests/unit/environments/dynamodb/types/test_unit_dynamodb_project_metadata.py +++ b/api/tests/unit/environments/dynamodb/types/test_unit_dynamodb_project_metadata.py @@ -2,6 +2,8 @@ from decimal import Decimal import pytest +from mypy_boto3_dynamodb.service_resource import Table +from pytest_mock import MockerFixture from environments.dynamodb.types import ( DynamoProjectMetadata, @@ -146,3 +148,29 @@ def test_finish_identity_migration_calls_put_item_with_correct_arguments( "triggered_at": None, } ) + + +def test_delete__removes_project_metadata_document_from_dynamodb( + flagsmith_project_metadata_table: Table, mocker: MockerFixture +): + # Given + first_project_id = 1 + mocker.patch( + "environments.dynamodb.types.project_metadata_table", + flagsmith_project_metadata_table, + ) + flagsmith_project_metadata_table.put_item(Item={"id": first_project_id}) + # Let's add another item to make sure that only the correct item is deleted + second_project_id = 2 + flagsmith_project_metadata_table.put_item(Item={"id": second_project_id}) + + project_metadata = DynamoProjectMetadata.get_or_new(first_project_id) + + # When + project_metadata.delete() + + # Then + assert flagsmith_project_metadata_table.scan()["Count"] == 1 + assert ( + flagsmith_project_metadata_table.scan()["Items"][0]["id"] == second_project_id + ) diff --git a/api/tests/unit/environments/dynamodb/wrappers/__init__.py b/api/tests/unit/environments/dynamodb/wrappers/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/api/tests/unit/environments/dynamodb/test_unit_dynamo_environment_wrapper.py b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamo_environment_wrapper.py similarity index 78% rename from api/tests/unit/environments/dynamodb/test_unit_dynamo_environment_wrapper.py rename to api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamo_environment_wrapper.py index 87a5d6f7eead..81d3993a6edb 100644 --- a/api/tests/unit/environments/dynamodb/test_unit_dynamo_environment_wrapper.py +++ b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamo_environment_wrapper.py @@ -1,5 +1,6 @@ import pytest from django.core.exceptions import ObjectDoesNotExist +from mypy_boto3_dynamodb.service_resource import Table from environments.dynamodb import DynamoEnvironmentWrapper from environments.models import Environment @@ -59,3 +60,19 @@ def test_get_item_raises_object_does_not_exists_if_get_item_does_not_return_any_ # Then with pytest.raises(ObjectDoesNotExist): dynamo_environment_wrapper.get_item(api_key) + + +def test_delete_environment__removes_environment_document_from_dynamodb( + dynamo_enabled_project_environment_one_document: dict, + dynamo_environment_wrapper: DynamoEnvironmentWrapper, + flagsmith_environment_table: Table, +): + # Given + api_key = dynamo_enabled_project_environment_one_document["api_key"] + assert flagsmith_environment_table.scan()["Count"] == 1 + + # When + dynamo_environment_wrapper.delete_environment(api_key) + + # Then + assert flagsmith_environment_table.scan()["Count"] == 0 diff --git a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_api_key_wrapper.py b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_api_key_wrapper.py similarity index 100% rename from api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_api_key_wrapper.py rename to api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_api_key_wrapper.py diff --git a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_v2_wrapper.py b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py similarity index 78% rename from api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_v2_wrapper.py rename to api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py index 805990e9a800..4e035f2d1f1b 100644 --- a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_v2_wrapper.py +++ b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py @@ -3,7 +3,7 @@ from mypy_boto3_dynamodb.service_resource import Table from pytest_django.fixtures import SettingsWrapper -from environments.dynamodb.dynamodb_wrapper import DynamoEnvironmentV2Wrapper +from environments.dynamodb import DynamoEnvironmentV2Wrapper from environments.dynamodb.types import ( IdentityOverridesV2Changeset, IdentityOverrideV2, @@ -154,3 +154,45 @@ def test_environment_v2_wrapper__write_environments__put_expected( results = flagsmith_environments_v2_table.scan()["Items"] assert len(results) == 1 assert results[0] == map_environment_to_environment_v2_document(environment) + + +def test_environment_v2_wrapper__delete_environment__deletes_related_data_from_dynamodb( + flagsmith_environments_v2_table: Table, + dynamodb_wrapper_v2: DynamoEnvironmentV2Wrapper, +) -> None: + # Given + environment_api_key = "api_key" + environment_id = "10" + + # Add some items to the table + for i in range(10): + flagsmith_environments_v2_table.put_item( + Item={ + "environment_api_key": environment_api_key, + "environment_id": environment_id, + "document_key": f"identity_override:{i}", + } + ) + + # Next, let's add an item for a different environment + environment_2_api_key = "different_api_key" + environment_2_id = "11" + flagsmith_environments_v2_table.put_item( + Item={ + "environment_api_key": environment_2_api_key, + "environment_id": environment_2_id, + "document_key": "identity_override:1", + } + ) + + # When + dynamodb_wrapper_v2.delete_environment(environment_id=environment_id) + + # Then + results = flagsmith_environments_v2_table.scan()["Items"] + assert len(results) == 1 + assert results[0] == { + "environment_api_key": environment_2_api_key, + "environment_id": environment_2_id, + "document_key": "identity_override:1", + } diff --git a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_identity_wrapper.py b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py similarity index 86% rename from api/tests/unit/environments/dynamodb/test_unit_dynamodb_identity_wrapper.py rename to api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py index 169ccd29295f..96ed023760dc 100644 --- a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_identity_wrapper.py +++ b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py @@ -6,6 +6,8 @@ from django.core.exceptions import ObjectDoesNotExist from flag_engine.identities.models import IdentityModel from flag_engine.segments.constants import IN +from mypy_boto3_dynamodb.service_resource import Table +from pytest_mock import MockerFixture from rest_framework.exceptions import NotFound from environments.dynamodb import DynamoIdentityWrapper @@ -18,8 +20,6 @@ ) if typing.TYPE_CHECKING: - from pytest_mock import MockerFixture - from environments.models import Environment from projects.models import Project @@ -229,11 +229,11 @@ def test_write_identities_skips_identity_if_identifier_is_too_large( def test_is_enabled_is_false_if_dynamo_table_name_is_not_set(settings, mocker): # Given mocker.patch( - "environments.dynamodb.dynamodb_wrapper.DynamoIdentityWrapper.table_name", + "environments.dynamodb.wrappers.identity_wrapper.DynamoIdentityWrapper.table_name", None, ) + mocked_boto3 = mocker.patch("environments.dynamodb.wrappers.base.boto3") - mocked_boto3 = mocker.patch("environments.dynamodb.dynamodb_wrapper.boto3") # When dynamo_identity_wrapper = DynamoIdentityWrapper() @@ -247,8 +247,8 @@ def test_is_enabled_is_true_if_dynamo_table_name_is_set(settings, mocker): # Given table_name = "random_table_name" settings.IDENTITIES_TABLE_NAME_DYNAMO = table_name - mocked_config = mocker.patch("environments.dynamodb.dynamodb_wrapper.Config") - mocked_boto3 = mocker.patch("environments.dynamodb.dynamodb_wrapper.boto3") + mocked_config = mocker.patch("environments.dynamodb.wrappers.base.Config") + mocked_boto3 = mocker.patch("environments.dynamodb.wrappers.base.boto3") # When dynamo_identity_wrapper = DynamoIdentityWrapper() @@ -276,7 +276,7 @@ def test_get_segment_ids_returns_correct_segment_ids( environment_document = map_environment_to_environment_document(environment) mocked_environment_wrapper = mocker.patch( - "environments.dynamodb.dynamodb_wrapper.DynamoEnvironmentWrapper" + "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper" ) mocked_environment_wrapper.return_value.get_item.return_value = environment_document @@ -321,7 +321,7 @@ def test_get_segment_ids_returns_segment_using_in_operator_for_integer_traits( environment_document = map_environment_to_environment_document(environment) mocked_environment_wrapper = mocker.patch( - "environments.dynamodb.dynamodb_wrapper.DynamoEnvironmentWrapper" + "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper" ) mocked_environment_wrapper.return_value.get_item.return_value = environment_document @@ -386,7 +386,7 @@ def test_get_segment_ids_with_identity_model(identity, environment, mocker): environment_document = map_environment_to_environment_document(environment) mocked_environment_wrapper = mocker.patch( - "environments.dynamodb.dynamodb_wrapper.DynamoEnvironmentWrapper" + "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper" ) mocked_environment_wrapper.return_value.get_item.return_value = environment_document @@ -408,12 +408,11 @@ def test_identity_wrapper__iter_all_items_paginated__returns_expected( environment_api_key = "test_api_key" limit = 1 - expected_engine_identity = IdentityModel.parse_obj(identity_document) expected_next_page_key = "next_page_key" environment_document = map_environment_to_environment_document(environment) mocked_environment_wrapper = mocker.patch( - "environments.dynamodb.dynamodb_wrapper.DynamoEnvironmentWrapper", + "environments.dynamodb.wrappers.environment_wrapper.DynamoEnvironmentWrapper", autospec=True, ) mocked_environment_wrapper.return_value.get_item.return_value = environment_document @@ -439,19 +438,59 @@ def test_identity_wrapper__iter_all_items_paginated__returns_expected( with pytest.raises(StopIteration): next(iterator) - assert result_1 == expected_engine_identity - assert result_2 == expected_engine_identity + assert result_1 == identity_document + assert result_2 == identity_document mocked_get_all_items.assert_has_calls( [ mocker.call( environment_api_key=environment_api_key, + projection_expression=None, limit=limit, ), mocker.call( environment_api_key=environment_api_key, limit=limit, + projection_expression=None, start_key=expected_next_page_key, ), ] ) + + +def test_delete_all_identities__deletes_all_identities_documents_from_dynamodb( + flagsmith_identities_table: Table, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + environment_api_key = "environment_one" + + # Let's create 2 identities for the same environment + identity_one = { + "composite_key": f"{environment_api_key}_identity_one", + "environment_api_key": environment_api_key, + "identifier": "identity_one", + } + identity_two = { + "composite_key": f"{environment_api_key}_identity_two", + "identifier": "identity_two", + "environment_api_key": environment_api_key, + } + + flagsmith_identities_table.put_item(Item=identity_one) + flagsmith_identities_table.put_item(Item=identity_two) + + # Let's create another identity for a different environment + identity_three = { + "composite_key": "environment_two_identity_one", + "identifier": "identity_three", + "environment_api_key": "environment_two", + } + flagsmith_identities_table.put_item(Item=identity_three) + + # When + dynamodb_identity_wrapper.delete_all_identities(environment_api_key) + + # Then + assert flagsmith_identities_table.scan()["Count"] == 1 + assert flagsmith_identities_table.scan()["Items"][0] == identity_three diff --git a/api/tests/unit/environments/test_unit_environments_models.py b/api/tests/unit/environments/test_unit_environments_models.py index 6ebf6077e2d5..d6b348fb3ed3 100644 --- a/api/tests/unit/environments/test_unit_environments_models.py +++ b/api/tests/unit/environments/test_unit_environments_models.py @@ -9,7 +9,9 @@ from core.request_origin import RequestOrigin from django.test import TestCase, override_settings from django.utils import timezone +from mypy_boto3_dynamodb.service_resource import Table from pytest_django.asserts import assertQuerysetEqual as assert_queryset_equal +from pytest_mock import MockerFixture from audit.models import AuditLog from audit.related_object_type import RelatedObjectType @@ -802,12 +804,16 @@ def test_get_hide_disabled_flags( assert environment.get_hide_disabled_flags() is expected_result -def test_saving_environment_api_key_calls_put_item_with_correct_arguments_if_enabled( - dynamo_enabled_project_environment_one, mocker +def test_saving_environment_api_key_creates_dynamo_document_if_enabled( + dynamo_enabled_project_environment_one: Environment, + mocker: MockerFixture, + flagsmith_environment_api_key_table: "Table", ): # Given - mocked_environment_api_key_wrapper = mocker.patch( - "environments.models.environment_api_key_wrapper", autospec=True + mocker.patch( + "environments.models.DynamoEnvironmentAPIKeyWrapper.table", + new_callable=mocker.PropertyMock, + return_value=flagsmith_environment_api_key_table, ) # When api_key = EnvironmentAPIKey.objects.create( @@ -815,7 +821,69 @@ def test_saving_environment_api_key_calls_put_item_with_correct_arguments_if_ena ) # Then - mocked_environment_api_key_wrapper.write_api_key.assert_called_with(api_key) + response = flagsmith_environment_api_key_table.get_item(Key={"key": api_key.key}) + assert response["Item"]["key"] == api_key.key + + +def test_deleting_environment_api_key_deletes_dynamo_document_if_enabled( + dynamo_enabled_project_environment_one: Environment, + mocker: MockerFixture, + flagsmith_environment_api_key_table: "Table", +): + # Given + mocker.patch( + "environments.models.DynamoEnvironmentAPIKeyWrapper.table", + new_callable=mocker.PropertyMock, + return_value=flagsmith_environment_api_key_table, + ) + api_key = EnvironmentAPIKey.objects.create( + name="Some key", environment=dynamo_enabled_project_environment_one + ) + assert flagsmith_environment_api_key_table.scan()["Count"] == 1 + + # When + api_key.delete() + + # Then + assert flagsmith_environment_api_key_table.scan()["Count"] == 0 + + +def test_deleting_environment_creates_task_to_delete_dynamo_document_if_enabled( + dynamo_enabled_project_environment_one: Environment, + mocker: MockerFixture, +) -> None: + # Given + mocked_task = mocker.patch("environments.tasks.delete_environment_from_dynamo") + mocker.patch( + "environments.models.DynamoEnvironmentWrapper.is_enabled", + new_callable=mocker.PropertyMock, + return_value=True, + ) + + # When + dynamo_enabled_project_environment_one.delete() + + # Then + mocked_task.delay.assert_called_once_with( + args=( + dynamo_enabled_project_environment_one.api_key, + dynamo_enabled_project_environment_one.id, + ) + ) + + +def test_delete_api_key_not_called_when_deleting_environment_api_key_for_non_edge_project( + environment_api_key: EnvironmentAPIKey, mocker: MockerFixture +) -> None: + # Given + mocked_environment_api_key_wrapper = mocker.patch( + "environments.models.environment_api_key_wrapper", autospec=True + ) + # When + environment_api_key.delete() + + # Then + mocked_environment_api_key_wrapper.delete_api_key.assert_not_called() def test_put_item_not_called_when_saving_environment_api_key_for_non_edge_project( diff --git a/api/tests/unit/environments/test_unit_environments_tasks.py b/api/tests/unit/environments/test_unit_environments_tasks.py index 6e454815a94c..0669c528176f 100644 --- a/api/tests/unit/environments/test_unit_environments_tasks.py +++ b/api/tests/unit/environments/test_unit_environments_tasks.py @@ -1,5 +1,8 @@ +from pytest_mock import MockerFixture + from audit.models import AuditLog from environments.tasks import ( + delete_environment_from_dynamo, process_environment_update, rebuild_environment_document, ) @@ -75,3 +78,33 @@ def test_process_environment_update_with_project_audit_log(environment, mocker): mock_send_environment_update_message_for_project.assert_called_once_with( environment.project ) + + +def test_delete_environment__calls_internal_methods_correctly( + mocker: MockerFixture, +) -> None: + # Given + environment_api_key = "test-api-key" + environment_id = 10 + + mocked_environment_wrapper = mocker.patch("environments.tasks.environment_wrapper") + mocked_environment_v2_wrapper = mocker.patch( + "environments.tasks.environment_v2_wrapper" + ) + DynamoIdentityWrapper = mocker.patch("environments.tasks.DynamoIdentityWrapper") + mocked_identity_wrapper = DynamoIdentityWrapper.return_value + + # When + delete_environment_from_dynamo(environment_api_key, environment_id) + + # Then + mocked_environment_wrapper.delete_environment.assert_called_once_with( + environment_api_key + ) + + mocked_environment_v2_wrapper.delete_environment.assert_called_once_with( + environment_id + ) + mocked_identity_wrapper.delete_all_identities.assert_called_once_with( + environment_api_key + ) diff --git a/api/tests/unit/features/test_unit_features_features_service.py b/api/tests/unit/features/test_unit_features_features_service.py index a02732840e5f..eba2ff70c4f5 100644 --- a/api/tests/unit/features/test_unit_features_features_service.py +++ b/api/tests/unit/features/test_unit_features_features_service.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from pytest_mock import MockerFixture - from environments.dynamodb.dynamodb_wrapper import ( + from environments.dynamodb import ( DynamoEnvironmentV2Wrapper, DynamoIdentityWrapper, )