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,
    TransactItems,
    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): ...

    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) as exc:
        dynamodb_persistence_layer.transact_write_items(transact)
    # print(exc.value)


def test_collection_get_item(
    dynamodb_seeds,
    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
):
    collect = DynamoDBCollection(dynamodb_persistence_layer)
    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_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
    assert 'ttl_date' 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('billing_policy', path_spec='payment_method')
        + SortKey('payment_policy'),
    )

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


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

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

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