"""Tests for util module."""

import json
from unittest.mock import Mock
from unittest.mock import patch

from flows.distribution_fee import util


@patch('flows.distribution_fee.util.boto3')
@patch('flows.distribution_fee.util.config')
def test_send_sns_notification(config_mock, boto_mock):
    """Test send_success_notification function."""
    config_mock.SNS_ACTION_SUCCESS = 'test passed'
    config_mock.SNS_SOURCE = 'py.test'
    config_mock.SNS_TOPIC_ARN = 'a unit test'

    sns = Mock()
    sns.publish.return_value = {'foo': 'bar'}
    boto_mock.client.return_value = sns

    correlation_id = '0123-4567-8910-1112'
    report_date = 'yesterday'
    upcs = ['123', '456']

    result = util.send_sns_notification(
        correlation_id, report_date, upcs)

    assert result == {'foo': 'bar'}

    expected_payload = {
        'action': config_mock.SNS_ACTION_SUCCESS,
        'correlation_id': correlation_id,
        'report_date': report_date,
        'source': config_mock.SNS_SOURCE,
        'upcs': upcs}

    params = sns.publish.call_args_list[0][1]

    assert params['TopicArn'] == config_mock.SNS_TOPIC_ARN
    assert params['Subject'] == config_mock.SNS_ACTION_SUCCESS
    assert json.loads(params['Message']) == expected_payload


@patch('flows.distribution_fee.util.datastore')
def test_get_etl_vendor_contracts(datastore, vendor_contracts):
    """Test get_etl_vendor_contracts function."""
    cursor = Mock()
    cursor.fetchall.return_value = vendor_contracts
    datastore.query.return_value = cursor

    expected_ids = {110, 220, 221, 330, 331, 332, 440, 550, 660}
    expected_map = {
        1: {10},
        2: {20, 21},
        3: {30, 31, 32},
        4: {40},
        5: {50},
        6: {60}}
    result_ids, result_map = util.get_etl_vendor_contracts('foo')

    # set cast for easier testing of unordered results
    assert set(result_ids) == expected_ids
    assert result_map == expected_map