diff --git a/aws_lambda_powertools/shared/functions.py b/aws_lambda_powertools/shared/functions.py index 16f51da1cb9..ea29ccf3ea5 100644 --- a/aws_lambda_powertools/shared/functions.py +++ b/aws_lambda_powertools/shared/functions.py @@ -140,8 +140,8 @@ def powertools_debug_is_set() -> bool: def slice_dictionary(data: dict, chunk_size: int) -> Generator[dict, None, None]: - for _ in range(0, len(data), chunk_size): - yield {dict_key: data[dict_key] for dict_key in itertools.islice(data, chunk_size)} + for i in range(0, len(data), chunk_size): + yield {key: data[key] for key in itertools.islice(data, i, i + chunk_size)} def extract_event_from_common_models(data: Any) -> dict | Any: diff --git a/tests/unit/test_shared_functions.py b/tests/unit/test_shared_functions.py index 2cd6a41aa12..1bf7c6e26a0 100644 --- a/tests/unit/test_shared_functions.py +++ b/tests/unit/test_shared_functions.py @@ -11,6 +11,7 @@ from aws_lambda_powertools.shared import constants from aws_lambda_powertools.shared.functions import ( abs_lambda_path, + slice_dictionary, extract_event_from_common_models, powertools_debug_is_set, powertools_dev_is_set, @@ -202,3 +203,17 @@ def test_sanitize_xray_segment_name_with_no_special_characters(): # THEN the sanitized name remains the same as the original name expected_name = valid_name assert sanitized_name == expected_name + + +@pytest.mark.parametrize( + "chunk_size, expected", + [ + (1, [{"k0": 0}, {"k1": 1}, {"k2": 2}, {"k3": 3}]), + (2, [{"k0": 0, "k1": 1}, {"k2": 2, "k3": 3}]), + (3, [{"k0": 0, "k1": 1, "k2": 2}, {"k3": 3}]), + ], +) +def test_slice_dictionary(chunk_size, expected): + data = {f"k{i}": i for i in range(4)} + chunks = list(slice_dictionary(data, chunk_size=chunk_size)) + assert chunks == expected