"""MSK module tests."""

import json
import base64
from unittest.mock import call, MagicMock, patch

import pytest
from lambdacommon.aws import msk


def test_msk_handler():
    """Test MSK handler."""
    msk_event = {
        'records': {
            'test_topic-2': [
                {
                    'topic': 'test_topic',
                    'partition': 2,
                    'offset': 0,
                    'timestamp': 1718198430933,
                    'timestampType': 'CREATE_TIME',
                    'key': 'ImYxM2FlZjZkLTVkMGQtNDIxOS1iYjc5LWFkMDljNzcyZTdmOSI=',
                    'value': base64.b64encode(json.dumps({'event_id': 1}).encode()),
                    'headers': []
                },
                {
                    'topic': 'test_topic',
                    'partition': 2,
                    'offset': 0,
                    'timestamp': 1718198430933,
                    'timestampType': 'CREATE_TIME',
                    'key': 'KvTxM2FlZjZkLTVkMGQtNDIxOS1iYjc5LWFkMDljNzcyZTdmOSI=',
                    'value': base64.b64encode(json.dumps({'event_id': 2}).encode()),
                    'headers': []
                }
            ]
        }
    }
    msk_context = MagicMock()

    mock_function = MagicMock()
    decorated_mock_function = msk.msk_handler()(mock_function)
    decorated_mock_function(msk_event, msk_context)

    assert mock_function.call_args_list == [
        call({'event_id': 1}, msk_context),
        call({'event_id': 2}, msk_context),
    ]


def test_msk_handler_process_exceptions():
    """Test MSK handler with exceptions processing."""
    msk_event = {
        'records': {
            'test_topic-2': [
                {
                    'topic': 'test_topic',
                    'partition': 2,
                    'offset': 0,
                    'timestamp': 1718198430933,
                    'timestampType': 'CREATE_TIME',
                    'key': 'ImYxM2FlZjZkLTVkMGQtNDIxOS1iYjc5LWFkMDljNzcyZTdmOSI=',
                    'value': base64.b64encode(json.dumps({'event_id': 1}).encode()),
                    'headers': []
                },
                {
                    'topic': 'test_topic',
                    'partition': 2,
                    'offset': 0,
                    'timestamp': 1718198430933,
                    'timestampType': 'CREATE_TIME',
                    'key': 'KvTxM2FlZjZkLTVkMGQtNDIxOS1iYjc5LWFkMDljNzcyZTdmOSI=',
                    'value': base64.b64encode(json.dumps({'event_id': 2}).encode()),
                    'headers': []
                }
            ]
        }
    }
    msk_context = MagicMock()

    mock_function = MagicMock()
    mock_function.side_effect = Exception
    decorated_mock_function = msk.msk_handler()(mock_function)

    # executes once and then fails on exception
    with pytest.raises(Exception):
        decorated_mock_function(msk_event, msk_context)
    assert mock_function.call_args_list == [
        call({'event_id': 1}, msk_context)
    ]

    mock_function.reset_mock()

    # swallows exceptions
    decorated_mock_function = msk.msk_handler(process_exceptions=True)(mock_function)
    decorated_mock_function(msk_event, msk_context)
    assert mock_function.call_args_list == [
        call({'event_id': 1}, msk_context),
        call({'event_id': 2}, msk_context)
    ]


@patch.object(msk, '_process_payload')
@patch.object(msk, 'ThreadPoolExecutor')
def test_msk_handler_concurrent(mock_executor, mock_process_payload):
    """Test MSK handler."""
    msk_event = {
        'records': {
            'test_topic-2': [
                {
                    'topic': 'test_topic',
                    'partition': 2,
                    'offset': 0,
                    'timestamp': 1718198430933,
                    'timestampType': 'CREATE_TIME',
                    'key': 'ImYxM2FlZjZkLTVkMGQtNDIxOS1iYjc5LWFkMDljNzcyZTdmOSI=',
                    'value': base64.b64encode(json.dumps({'event_id': 1}).encode()),
                    'headers': []
                },
                {
                    'topic': 'test_topic',
                    'partition': 2,
                    'offset': 0,
                    'timestamp': 1718198430933,
                    'timestampType': 'CREATE_TIME',
                    'key': 'KvTxM2FlZjZkLTVkMGQtNDIxOS1iYjc5LWFkMDljNzcyZTdmOSI=',
                    'value': base64.b64encode(json.dumps({'event_id': 2}).encode()),
                    'headers': []
                }
            ]
        }
    }
    msk_context = MagicMock()

    mock_function = MagicMock()
    decorated_mock_function = msk.msk_handler(concurrent=True)(mock_function)
    decorated_mock_function(msk_event, msk_context)

    assert mock_executor.return_value.__enter__.return_value.submit.call_args_list == [
        call(mock_process_payload, mock_function, {'event_id': 1}, msk_context, False),
        call(mock_process_payload, mock_function, {'event_id': 2}, msk_context, False),
    ]