Skip to content

Commit

Permalink
fix: No pagination when querying environments_v2 (#3661)
Browse files Browse the repository at this point in the history
  • Loading branch information
@khvn26
khvn26 committed Mar 25, 2024
1 parent a69ff43 commit 7e19f4f
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 15 deletions.
19 changes: 10 additions & 9 deletions api/environments/dynamodb/wrappers/environment_wrapper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,17 +67,18 @@ def get_identity_overrides_by_environment_id(
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 list(
self.query_get_all_items(
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

Expand Down
67 changes: 61 additions & 6 deletions api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,16 @@

from mypy_boto3_dynamodb.service_resource import Table
from pytest_django.fixtures import SettingsWrapper
from pytest_mock import MockerFixture

from environments.dynamodb import DynamoEnvironmentV2Wrapper
from environments.dynamodb.types import (
IdentityOverridesV2Changeset,
IdentityOverrideV2,
)
from environments.dynamodb.utils import (
get_environments_v2_identity_override_document_key,
)
from environments.models import Environment
from features.models import Feature, FeatureState
from util.mappers import (
Expand All @@ -31,7 +35,9 @@ def test_environment_v2_wrapper__get_identity_overrides_by_environment_id__retur
identifier = "identity1"
override_document = {
"environment_id": str(environment.id),
"document_key": f"identity_override:{feature.id}:{identity_uuid}",
"document_key": get_environments_v2_identity_override_document_key(
feature_id=feature.id, identity_uuid=identity_uuid
),
"environment_api_key": environment.api_key,
"identifier": identifier,
"feature_state": {},
Expand All @@ -53,6 +59,45 @@ def test_environment_v2_wrapper__get_identity_overrides_by_environment_id__retur
assert results[0] == override_document


def test_environment_v2_wrapper__get_identity_overrides_by_environment_id__last_evaluated_key__call_expected(
flagsmith_environments_v2_table: Table,
mocker: MockerFixture,
) -> None:
# Given
wrapper = DynamoEnvironmentV2Wrapper()
mocker.patch.object(
wrapper, "get_table"
).return_value = table_mock = mocker.MagicMock(spec=flagsmith_environments_v2_table)

last_evaluated_key = "next_page_key"
override_document = {"test": "document"}

table_mock.query.side_effect = [
{"Items": [override_document], "LastEvaluatedKey": last_evaluated_key},
{"Items": [override_document], "LastEvaluatedKey": None},
]

# When
results = wrapper.get_identity_overrides_by_environment_id(
environment_id=mocker.ANY,
feature_id=mocker.ANY,
)

# Then
assert results == [override_document, override_document]
wrapper.table.query.assert_has_calls(
[
mocker.call(
KeyConditionExpression=mocker.ANY,
),
mocker.call(
KeyConditionExpression=mocker.ANY,
ExclusiveStartKey=last_evaluated_key,
),
]
)


def test_environment_v2_wrapper__update_identity_overrides__put_expected(
settings: SettingsWrapper,
environment: Environment,
Expand All @@ -69,7 +114,9 @@ def test_environment_v2_wrapper__update_identity_overrides__put_expected(
override_document = IdentityOverrideV2.parse_obj(
{
"environment_id": str(environment.id),
"document_key": f"identity_override:{feature.id}:{identity_uuid}",
"document_key": get_environments_v2_identity_override_document_key(
feature_id=feature.id, identity_uuid=identity_uuid
),
"environment_api_key": environment.api_key,
"feature_state": map_feature_state_to_engine(feature_state),
"identifier": identifier,
Expand Down Expand Up @@ -110,7 +157,9 @@ def test_environment_v2_wrapper__update_identity_overrides__delete_expected(
IdentityOverrideV2.parse_obj(
{
"environment_id": str(environment.id),
"document_key": f"identity_override:{feature.id}:{identity_uuid}",
"document_key": get_environments_v2_identity_override_document_key(
feature_id=feature.id, identity_uuid=identity_uuid
),
"environment_api_key": environment.api_key,
"feature_state": map_feature_state_to_engine(feature_state),
"identifier": identifier,
Expand Down Expand Up @@ -170,7 +219,9 @@ def test_environment_v2_wrapper__delete_environment__deletes_related_data_from_d
Item={
"environment_api_key": environment_api_key,
"environment_id": environment_id,
"document_key": f"identity_override:{i}",
"document_key": get_environments_v2_identity_override_document_key(
feature_id=i,
),
}
)

Expand All @@ -181,7 +232,9 @@ def test_environment_v2_wrapper__delete_environment__deletes_related_data_from_d
Item={
"environment_api_key": environment_2_api_key,
"environment_id": environment_2_id,
"document_key": "identity_override:1",
"document_key": get_environments_v2_identity_override_document_key(
feature_id=1,
),
}
)

Expand All @@ -194,5 +247,7 @@ def test_environment_v2_wrapper__delete_environment__deletes_related_data_from_d
assert results[0] == {
"environment_api_key": environment_2_api_key,
"environment_id": environment_2_id,
"document_key": "identity_override:1",
"document_key": get_environments_v2_identity_override_document_key(
feature_id=1,
),
}

0 comments on commit 7e19f4f

Please sign in to comment.