"""Test dynamodb utils."""
import types
from unittest.mock import MagicMock

from feed_ingestion.util.aws import dynamodb


def test_query_dynamodb_paginated_generator():
    query_args = {
        'testKey': 'testValue',
    }
    dynamodb_table_mock = MagicMock()
    dynamodb_table_mock.query.side_effect = [
        {
            'Items': [11, 12, 13],
            'LastEvaluatedKey': 222,
        },
        {
            'Items': [21, 22],
        },
    ]
    result = dynamodb.query_dynamodb_paginated_generator(
        dynamodb_table=dynamodb_table_mock,
        query_args=query_args
    )
    assert isinstance(result, types.GeneratorType)
    result = list(result)
    assert result == [11, 12, 13, 21, 22]