saladeaula.digital/layercake/tests/test_dynamodb.py

from datetime import datetime
from decimal import Decimal
from ipaddress import IPv4Address

import pytest
from layercake.dateutils import ttl
from layercake.dynamodb import (
    ComposeKey,
    DynamoDBCollection,
    DynamoDBPersistenceLayer,
    KeyPair,
    PartitionKey,
    PrefixKey,
    SortKey,
    TransactKey,
    serialize,
)


def test_serialize():
    assert serialize(
        {
            'id': '123',
            'sk': 'abc',
            'date': datetime.fromisoformat('2025-03-20T18:29:10.713994'),
            'ip': IPv4Address('127.0.0.1'),
        }
    ) == {
        'id': {'S': '123'},
        'sk': {'S': 'abc'},
        'date': {'S': '2025-03-20T18:29:10.713994'},
        'ip': {'S': '127.0.0.1'},
    }

    assert serialize(
        {'ids': ('1', '2', '3')},
    ) == {
        'ids': {
            'L': [{'S': '1'}, {'S': '2'}, {'S': '3'}],
        }
    }

    assert serialize(
        {'ids': ['1', '2', '3']},
    ) == {
        'ids': {
            'L': [{'S': '1'}, {'S': '2'}, {'S': '3'}],
        }
    }

    assert serialize({'ids': {'1'}}) == {'ids': {'SS': ['1']}}


def test_composekey():
    key = ComposeKey(('122', 'abc'), prefix='schedules', delimiter=':')
    assert key == 'schedules:122:abc'
    assert key.prefix == 'schedules'
    assert key.delimiter == ':'

    assert ComposeKey(('122', 'abc')) == '122#abc'
    assert ComposeKey('122') == '122'


def test_partitionkey():
    assert PartitionKey('123') == {'id': '123'}
    assert PartitionKey('123').expr_attr_name() == {'#pk': 'id'}
    assert PartitionKey('123').expr_attr_values() == {':pk': '123'}


def test_keypair():
    assert KeyPair('123', 'abc') == {'id': '123', 'sk': 'abc'}
    assert KeyPair('123', 'abc').expr_attr_name() == {'#pk': 'id', '#sk': 'sk'}
    assert KeyPair('123', 'abc').expr_attr_values() == {':pk': '123', ':sk': 'abc'}

    assert KeyPair.parse_obj({'id': '123', 'sk': 'abc'}) == {'id': '123', 'sk': 'abc'}
    assert KeyPair.parse_obj({'sk': 'abc', 'id': '123'}) == {'id': '123', 'sk': 'abc'}
    assert KeyPair.parse_obj(['123', 'abc']) == {'id': '123', 'sk': 'abc'}
    assert KeyPair.parse_obj([]) is None


def test_prefixkey():
    key = PrefixKey('emails')
    assert key == 'emails#'
    assert isinstance(key, PrefixKey)

    delimiter = PrefixKey('emails', None)
    assert delimiter == 'emails'


def test_transact_write_items(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    class EmailConflictError(Exception): ...

    with dynamodb_persistence_layer.transact_items() as transact:
        # transact = TransactItems(dynamodb_persistence_layer.table_name)
        transact.put(item=KeyPair('5OxmMjL-ujoR5IMGegQz', '0'))
        transact.put(item=KeyPair('cpf', '07879819908'))
        transact.put(
            item=KeyPair('email', 'sergio@somosbeta.com.br'),
            cond_expr='attribute_not_exists(sk)',
        )
        transact.put(
            item=KeyPair(
                '5OxmMjL-ujoR5IMGegQz',
                ComposeKey('sergio@somosbeta.com.br', 'emails'),
            ),
            cond_expr='attribute_not_exists(sk)',
            exc_cls=EmailConflictError,
        )

        with pytest.raises(EmailConflictError):
            transact.write_items()


def test_collection_get_item(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = dynamodb_persistence_layer.collect
    data_notfound = collect.get_item(
        KeyPair(
            pk='5OxmMjL-ujoR5IMGegQz',
            sk='tenant',
        ),
        raise_on_error=False,
        default={},
    )
    assert data_notfound == {}

    # This data was added from seeds
    data = collect.get_item(
        KeyPair(
            pk='5OxmMjL-ujoR5IMGegQz',
            sk=ComposeKey('sergio@somosbeta.com.br', prefix='emails'),
        ),
        default={},
    )
    assert data == {
        'email_verified': True,
        'mx_record_exists': True,
        'sk': 'emails#sergio@somosbeta.com.br',
        'email_primary': True,
        'id': '5OxmMjL-ujoR5IMGegQz',
        'create_date': '2019-03-25T00:00:00-03:00',
        'update_date': '2023-11-09T12:13:04.308986-03:00',
    }

    class NotFoundError(Exception): ...

    with pytest.raises(NotFoundError):
        collect.get_item(
            KeyPair('5OxmMjL-ujoR5IMGegQz', 'notfound'),
            exc_cls=NotFoundError,
        )


def test_collection_get_item_path_spec(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = dynamodb_persistence_layer.collect

    # This data was added from seeds
    data = collect.get_item(
        KeyPair(
            pk='5OxmMjL-ujoR5IMGegQz',
            sk=SortKey(
                ComposeKey('sergio@somosbeta.com.br', prefix='emails'),
                path_spec='mx_record_exists',
            ),
        ),
        default={},
    )
    assert data


def test_collection_put_item(
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)

    assert collect.put_item(
        KeyPair(
            '5OxmMjL-ujoR5IMGegQz',
            ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
        ),
        name='Beta Educação',
        ttl=ttl(days=3),
    )

    data = collect.get_item(
        KeyPair(
            pk='5OxmMjL-ujoR5IMGegQz',
            sk=ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
        ),
    )

    assert data['sk'] == 'orgs#6d1044d5-18c5-437c-9219-fc2ace7e5ebc'
    assert 'name' in data
    assert 'ttl' in data


def test_collection_delete_item(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)

    # This data was added from seeds
    assert collect.delete_item(
        KeyPair(
            '5OxmMjL-ujoR5IMGegQz',
            ComposeKey('sergio@somsbeta.com.br', prefix='emails'),
        )
    )


def test_collection_query(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)

    # This data was added from seeds
    logs = collect.query(
        PartitionKey(
            ComposeKey('5OxmMjL-ujoR5IMGegQz', prefix='logs'),
        ),
    )
    assert len(logs['items']) == 2
    assert logs == {
        'items': [
            {
                'sk': '2024-02-08T16:42:33.776409-03:00',
                'action': 'OPEN_EMAIL',
                'id': '5OxmMjL-ujoR5IMGegQz',
            },
            {
                'sk': '2019-03-25T00:00:00-03:00',
                'action': 'CLICK_EMAIL',
                'id': '5OxmMjL-ujoR5IMGegQz',
            },
        ],
        'last_key': None,
    }

    # This data was added from seeds
    emails = collect.query(
        KeyPair('5OxmMjL-ujoR5IMGegQz', PrefixKey('emails')),
    )
    assert emails == {
        'items': [
            {
                'email_verified': True,
                'mx_record_exists': True,
                'sk': 'sergio@somosbeta.com.br',  # Removed prefix from Sort Key
                'email_primary': True,
                'id': '5OxmMjL-ujoR5IMGegQz',
                'create_date': '2019-03-25T00:00:00-03:00',
                'update_date': '2023-11-09T12:13:04.308986-03:00',
            }
        ],
        'last_key': None,
    }


def test_collection_get_items(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)
    doc = collect.get_items(
        TransactKey('cJtK9SsnJhKPyxESe7g3DG')
        + SortKey('0')
        + SortKey('metadata#billing_policy', path_spec='payment_method')
        + SortKey('metadata#payment_policy', remove_prefix='metadata#'),
    )

    assert doc == {
        'sk': '0',
        'name': 'EDUSEG',
        'id': 'cJtK9SsnJhKPyxESe7g3DG',
        'cnpj': '15608435000190',
        'email': 'org+15608435000190@users.noreply.betaeducacao.com.br',
        'metadata#billing_policy': 'PIX',
        'payment_policy': {'due_days': Decimal('90')},
    }


def test_collection_get_items_not_found(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)
    doc = collect.get_items(
        TransactKey('not_found')
        + SortKey('0')
        + SortKey('metadata#not_found', path_spec='payment_method')
    )

    assert doc == {}


def test_collection_get_items_unflatten(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)

    doc = collect.get_items(
        TransactKey('cJtK9SsnJhKPyxESe7g3DG')
        + SortKey('metadata#billing_policy')
        + SortKey('metadata#payment_policy', remove_prefix='metadata#'),
        flatten_top=False,
    )

    assert doc == {
        'metadata#billing_policy': {
            'billing_day': Decimal('1'),
            'payment_method': 'PIX',
        },
        'payment_policy': {'due_days': Decimal('90')},
    }


def test_collection_get_items_pair(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)
    doc = collect.get_items(
        KeyPair('5OxmMjL-ujoR5IMGegQz', '0')
        + KeyPair('cpf', '07879819908')
        + KeyPair('email', 'osergiosiqueira@gmail.com')
    )

    assert doc == {
        'tenant:org_id': [
            'cJtK9SsnJhKPyxESe7g3DG',
            'edp8njvgQuzNkLx2ySNfAD',
            '8TVSi5oACLxTiT8ycKPmaQ',
        ],
        'email_verified': True,
        'last_login': '2024-02-08T20:53:45.818126-03:00',
        'sk': '0',
        'cpf': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
        'name': 'Sérgio Rafael de Siqueira',
        'id': '5OxmMjL-ujoR5IMGegQz',
        'create_date': '2019-03-25T00:00:00-03:00',
        'cognito:sub': '58efed8d-d276-41a8-8502-4ab8b5a6415e',
        'update_date': '2024-02-08T16:42:33.776409-03:00',
        'email': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
    }


def test_collection_get_items_pair_unflatten(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)
    doc = collect.get_items(
        KeyPair('5OxmMjL-ujoR5IMGegQz', '0')
        + KeyPair('cpf', '07879819908')
        + KeyPair('email', 'osergiosiqueira@gmail.com'),
        flatten_top=False,
    )

    assert doc == {
        '5OxmMjL-ujoR5IMGegQz': {
            'tenant:org_id': [
                'cJtK9SsnJhKPyxESe7g3DG',
                'edp8njvgQuzNkLx2ySNfAD',
                '8TVSi5oACLxTiT8ycKPmaQ',
            ],
            'email_verified': True,
            'last_login': '2024-02-08T20:53:45.818126-03:00',
            'cpf': '07879819908',
            'name': 'Sérgio Rafael de Siqueira',
            'create_date': '2019-03-25T00:00:00-03:00',
            'cognito:sub': '58efed8d-d276-41a8-8502-4ab8b5a6415e',
            'update_date': '2024-02-08T16:42:33.776409-03:00',
            'email': 'sergio@somosbeta.com.br',
        },
        'cpf': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
        'email': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
    }


def test_collection_get_items_pair_path_spec(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)
    doc = collect.get_items(
        KeyPair(
            'cpf',
            SortKey('07879819908', path_spec='user_id'),
            rename_key='user_id',
        )
        + KeyPair(
            'email',
            SortKey('osergiosiqueira@gmail.com', path_spec='user_id'),
            retain_key=True,
        ),
        flatten_top=False,
    )

    assert doc == {
        'user_id': '5OxmMjL-ujoR5IMGegQz',
        'email': '5OxmMjL-ujoR5IMGegQz',
    }