"""Test cases for move_statement_attachments per label dir script."""
import copy

import pytest

from unittest import mock

from botocore import exceptions

import move_attachments_per_label_dir as script


@pytest.mark.parametrize('key, expected_result', [
    ('L123_123', True),
    ('L123_123_123_123', False)
])
def test_is_month_item(key, expected_result):
    """Test is_month_item utility function."""
    item = {'label_type_id_period_id': key}
    assert script.is_month_item(item) == expected_result


@mock.patch('move_attachments_per_label_dir.util')
def test_get_dest_file_key(util):
    """Test get_dest_file_key utility function."""
    label_type_id = 'L123'
    file_key = 'test/me.key'
    full_file_key = 'me.key'
    dest_dir = 'test-test'
    util.get_file_name_from_key.return_value = full_file_key
    expected_result = 'test-test/L123/me.key'

    result = script.get_dest_file_key(label_type_id, file_key, dest_dir)
    assert result == expected_result
    util.get_file_name_from_key.assert_called_once_with(file_key)


@mock.patch('move_attachments_per_label_dir.dynamodb')
def test_get_approximate_table_count(dynamodb):
    """Test get_approximate_table_count utility function."""
    script.TABLE_NAME = 'test-{env}-test'
    expected_table_name = 'test-devdev-test'
    expected_result = 42
    mock_table = mock.MagicMock(item_count=expected_result)
    dynamodb.get_dynamodb_table.return_value = mock_table
    assert script.get_approximate_table_count('devdev') == expected_result
    dynamodb.get_dynamodb_table.assert_called_once_with(expected_table_name)


@pytest.mark.parametrize('item_count, limit, expected_result', [
    (100, 100, '0:00:02'),
    (100, 10, '0:00:20'),
    (0, 100, '0:00:00')
])
def test_get_estimated_time(item_count, limit, expected_result):
    """Test get_estimated_time utility function."""
    result = script.get_estimated_time(item_count, limit)
    assert result == expected_result


@pytest.mark.parametrize('dynamodb_key, s3_key', [
    ('key1', 's3/key'),
    ('key2', 'some_other_key'),
])
def test_build_update_kwargs(dynamodb_key, s3_key):
    """Test build_update_kwargs utility function."""
    expected_result = {
        'Key': dynamodb_key,
        'UpdateExpression': 'SET file_key = :file_key',
        'ExpressionAttributeValues': {':file_key': s3_key},
    }
    result = script.build_update_kwargs(dynamodb_key, s3_key)
    assert result == expected_result


@mock.patch('move_attachments_per_label_dir.util')
def test_get_primary_keys(util):
    """Test get_primary_keys utility function."""
    key_month = 'L132_321'
    file_name = 'test_file.name'
    quarter_periods = [320, 321, 322]
    util.get_quarter_periods_by_period.return_value = quarter_periods
    quarter_key = 'L123_320_321_322'
    util.get_attachment_primary_key_value.return_value = quarter_key
    quarter_file_name = 'quarter_file_name'
    util.format_quarter_file_name.return_value = quarter_file_name
    expected_result = [
        {
            'label_type_id_period_id': key_month,
            'file_name': file_name},
        {
            'label_type_id_period_id': quarter_key,
            'file_name': quarter_file_name}
    ]

    result = script.get_primary_keys(key_month, file_name)
    assert result == expected_result


@mock.patch('move_attachments_per_label_dir.build_update_kwargs')
@mock.patch('move_attachments_per_label_dir.dynamodb')
@mock.patch('move_attachments_per_label_dir.s3')
def test_move_attachment_and_update_record_successful(
        s3, dynamodb, build_update_kwargs):
    """Test move_attachment_and_update_record utility function."""
    env = 'testenv'
    script.TABLE_NAME = 'test-{env}'
    expected_table_name = 'test-testenv'
    bucket = 'bucket-test'
    source_key = 'src/key'
    dest_key = 'dst/key'
    dynamodb_keys = ['one', 'two']

    s3.copy_object.return_value = True

    mock_table = mock.MagicMock()
    dynamodb.get_dynamodb_table.return_value = mock_table

    update_kwargs = ['update_one', 'update_two']
    build_update_kwargs.side_effect = update_kwargs

    dynamodb.update_existing_item.return_value = True

    assert script.move_attachment_and_update_record(
        env, bucket, source_key, dest_key, dynamodb_keys)
    s3.copy_object.assert_called_once_with(
        bucket=bucket,
        source_key=source_key,
        dest_key=dest_key)
    dynamodb.get_dynamodb_table.assert_called_once_with(expected_table_name)
    build_update_kwargs.assert_has_calls([
        mock.call(dynamodb_key, dest_key)
        for dynamodb_key in dynamodb_keys])
    dynamodb.update_existing_item.assert_has_calls([
        mock.call(mock_table, update_data) for update_data in update_kwargs])
    s3.delete_object.assert_called_once_with(bucket, source_key)


@mock.patch('move_attachments_per_label_dir.build_update_kwargs')
@mock.patch('move_attachments_per_label_dir.dynamodb')
@mock.patch('move_attachments_per_label_dir.s3')
def test_move_attachment_and_update_record_error_on_copy(
        s3, dynamodb, build_update_kwargs):
    """Test move_attachment_and_update_record utility function."""
    env = 'testenv'
    script.TABLE_NAME = 'test-{env}'
    bucket = 'bucket-test'
    source_key = 'src/key'
    dest_key = 'dst/key'
    dynamodb_keys = ['one', 'two']

    s3.copy_object.return_value = False

    mock_table = mock.MagicMock()
    dynamodb.get_dynamodb_table.return_value = mock_table

    update_kwargs = ['update_one', 'update_two']
    build_update_kwargs.side_effect = update_kwargs

    dynamodb.update_existing_item.return_value = True

    assert not script.move_attachment_and_update_record(
        env, bucket, source_key, dest_key, dynamodb_keys)
    s3.copy_object.assert_called_once_with(
        bucket=bucket,
        source_key=source_key,
        dest_key=dest_key)
    assert dynamodb.get_dynamodb_table.call_count == 0
    assert build_update_kwargs.call_count == 0
    assert dynamodb.update_existing_item.call_count == 0
    assert s3.delete_object.call_count == 0


@mock.patch('move_attachments_per_label_dir.build_update_kwargs')
@mock.patch('move_attachments_per_label_dir.dynamodb')
@mock.patch('move_attachments_per_label_dir.s3')
def test_move_attachment_and_update_record_error_on_update(
        s3, dynamodb, build_update_kwargs):
    """Test move_attachment_and_update_record utility function."""
    env = 'testenv'
    script.TABLE_NAME = 'test-{env}'
    expected_table_name = 'test-testenv'
    bucket = 'bucket-test'
    source_key = 'src/key'
    dest_key = 'dst/key'
    dynamodb_keys = ['one', 'two']

    s3.copy_object.return_value = True

    mock_table = mock.MagicMock()
    dynamodb.get_dynamodb_table.return_value = mock_table

    update_kwargs = ['update_one', 'update_two']
    build_update_kwargs.side_effect = update_kwargs

    dynamodb.update_existing_item.side_effect = exceptions.ClientError({}, '')

    assert script.move_attachment_and_update_record(
        env, bucket, source_key, dest_key, dynamodb_keys)
    s3.copy_object.assert_called_once_with(
        bucket=bucket,
        source_key=source_key,
        dest_key=dest_key)
    dynamodb.get_dynamodb_table.assert_called_once_with(expected_table_name)
    build_update_kwargs.assert_has_calls([
        mock.call(dynamodb_key, dest_key)
        for dynamodb_key in dynamodb_keys])
    dynamodb.update_existing_item.assert_has_calls([
        mock.call(mock_table, update_data) for update_data in update_kwargs])
    assert s3.delete_object.call_count == 0


@mock.patch(
    'move_attachments_per_label_dir.move_attachment_and_update_record')
@mock.patch('move_attachments_per_label_dir.get_dest_file_key')
@mock.patch('move_attachments_per_label_dir.get_primary_keys')
def test_process_items(
        get_primary_keys, get_dest_file_key,
        move_attachment_and_update_record):
    """Test process_items utility function."""
    env = 'testenv'
    bucket = 'test-bucket'
    source_key = 'file_key'
    file_name = 'file_name'
    key_month = 'L123_321'
    key_quarter = 'L123_321_322_323'
    dest_key = 'dest_key'
    dest_dir = 'test-dest-dir'

    get_dest_file_key.return_value = dest_key
    item_month = {
        'label_type_id_period_id': key_month,
        'file_name': file_name,
        'bucket_name': bucket,
        'file_key': source_key,
    }
    item_quarter = {
        'label_type_id_period_id': key_quarter,
        'file_name': file_name,
        'bucket_name': bucket,
        'file_key': source_key,
    }
    items = [item_month, item_quarter]
    keys_to_update = ['key1', 'key2']
    get_primary_keys.return_value = keys_to_update

    expected_result = {
        'successful_count': 1,
        'errors': []
    }
    result = script.process_items(env, items, dest_dir)
    assert result == expected_result
    move_attachment_and_update_record.assert_called_once_with(
        env=env,
        bucket=bucket,
        source_key=source_key,
        dest_key=dest_key,
        dynamodb_keys=keys_to_update)
    get_dest_file_key.assert_called_once_with('L123', source_key, dest_dir)
    get_primary_keys.assert_called_once_with(key_month, file_name)


@mock.patch(
    'move_attachments_per_label_dir.move_attachment_and_update_record')
@mock.patch('move_attachments_per_label_dir.get_dest_file_key')
@mock.patch('move_attachments_per_label_dir.get_primary_keys')
def test_process_items_quarter(
        get_primary_keys, get_dest_file_key,
        move_attachment_and_update_record):
    """Test process_items utility function."""
    env = 'testenv'
    bucket = 'test-bucket'
    source_key = 'file_key'
    file_name = 'file_name'
    dynamodb_key = 'L123_321_322_323'
    dest_key = 'dest_key'
    dest_dir = 'test-dest-dir'

    get_dest_file_key.return_value = dest_key
    item = {
        'label_type_id_period_id': dynamodb_key,
        'file_name': file_name,
        'bucket_name': bucket,
        'file_key': source_key,
    }
    items = [item]
    keys_to_update = ['key1', 'key2']
    get_primary_keys.return_value = keys_to_update

    expected_result = {
        'successful_count': 0,
        'errors': []
    }
    result = script.process_items(env, items, dest_dir)
    assert result == expected_result
    assert move_attachment_and_update_record.call_count == 0
    assert get_dest_file_key.call_count == 0
    assert get_primary_keys.call_count == 0


@mock.patch(
    'move_attachments_per_label_dir.move_attachment_and_update_record')
@mock.patch('move_attachments_per_label_dir.get_dest_file_key')
@mock.patch('move_attachments_per_label_dir.get_primary_keys')
def test_process_items_error(
        get_primary_keys, get_dest_file_key,
        move_attachment_and_update_record):
    """Test process_items utility function."""
    env = 'testenv'
    bucket = 'test-bucket'
    source_key = 'file_key'
    file_name = 'file_name'
    key_month = 'L123_321'
    key_quarter = 'L123_321_322_323'
    dest_key = 'dest_key'
    dest_dir = 'test-dest-dir'

    get_dest_file_key.return_value = dest_key
    item_month = {
        'label_type_id_period_id': key_month,
        'file_name': file_name,
        'bucket_name': bucket,
        'file_key': source_key,
    }
    item_quarter = {
        'label_type_id_period_id': key_quarter,
        'file_name': file_name,
        'bucket_name': bucket,
        'file_key': source_key,
    }
    items = [item_month, item_quarter]
    keys_to_update = ['key1', 'key2']
    get_primary_keys.return_value = keys_to_update

    expected_result = {
        'successful_count': 0,
        'errors': [copy.deepcopy(item_month)]
    }
    move_attachment_and_update_record.return_value = False
    result = script.process_items(env, items, dest_dir)
    assert result == expected_result
    move_attachment_and_update_record.assert_called_once_with(
        env=env,
        bucket=bucket,
        source_key=source_key,
        dest_key=dest_key,
        dynamodb_keys=keys_to_update)
    get_dest_file_key.assert_called_once_with('L123', source_key, dest_dir)
    get_primary_keys.assert_called_once_with(key_month, file_name)


@mock.patch('move_attachments_per_label_dir.common_config')
@mock.patch('move_attachments_per_label_dir.conditions')
@mock.patch('move_attachments_per_label_dir.dynamodb')
def test_get_data_from_dynamodb(dynamodb, conditions, common_config):
    """Test get_data_from_dynamodb utility function."""
    env = 'testenv'
    expected_attrs_to_get = 'a,b,v'
    script.ATTRS_TO_GET = expected_attrs_to_get
    script.TABLE_NAME = 'test-{env}'
    expected_table_name = 'test-testenv'
    source = 'test/'
    limit = 4242

    table = mock.MagicMock()
    dynamodb.get_dynamodb_table.return_value = table

    expected_scan_kwargs = {
        'Select': 'SPECIFIC_ATTRIBUTES',
        'ProjectionExpression': expected_attrs_to_get,
        'FilterExpression': mock.ANY,
        'ConsistentRead': True,
        'Limit': limit,
        'ReturnConsumedCapacity': 'TOTAL'
    }
    dynamodb.get_scan_results.return_value = ({'a': 'a'} for _ in range(3))
    assert all(script.get_data_from_dynamodb(env, source, limit))

    dynamodb.get_dynamodb_table.assert_called_once_with(expected_table_name)
    conditions.Attr.assert_has_calls([
        mock.call('file_key').begins_with(source)])
    dynamodb.get_scan_results.assert_called_once_with(
        table, expected_scan_kwargs)


@mock.patch('move_attachments_per_label_dir.process_items')
@mock.patch('move_attachments_per_label_dir.get_data_from_dynamodb')
def test_move_valid_statement_attachments(
        get_data_from_dynamodb, process_items):
    """Test move_valid_statement_attachments utility function."""
    env = 'testenv'
    source_dir = 'test/'
    dest_dir = 'per-test'
    limit = 4242
    expected_items = 'items'
    item = {
        'Items': expected_items,
        'Count': 1,
        'ScannedCount': 42,
    }
    chunks = [item]
    get_data_from_dynamodb.return_value = chunks
    expected_result = {
        'Count': 1,
        'ScannedCount': 42,
        'successful_count': 1,
        'errors': []
    }
    process_items.return_value = {
        'successful_count': 1,
        'errors': []
    }
    result = script.move_valid_statement_attachments(
        env, source_dir, dest_dir, limit)
    assert result == expected_result
    get_data_from_dynamodb.assert_called_once_with(
        env, source_dir, limit)
    process_items.assert_called_once_with(
        env=env, items=expected_items, dest_dir=dest_dir)


@mock.patch('move_attachments_per_label_dir.process_items')
@mock.patch('move_attachments_per_label_dir.get_data_from_dynamodb')
def test_move_valid_statement_attachments_with_error(
        get_data_from_dynamodb, process_items):
    """Test move_valid_statement_attachments utility function."""
    env = 'testenv'
    source_dir = 'test/'
    dest_dir = 'per-test'
    limit = 4242
    expected_items = 'items'
    expected_items2 = 'items2'
    item = {
        'Items': expected_items,
        'Count': 1,
        'ScannedCount': 42,
    }
    item2 = {
        'Items': expected_items2,
        'Count': 1,
        'ScannedCount': 42,
    }
    chunks = [item, item2]
    get_data_from_dynamodb.return_value = chunks
    expected_result = {
        'Count': 2,
        'ScannedCount': 84,
        'successful_count': 1,
        'errors': [item2]
    }
    process_items.side_effect = [
        {
            'successful_count': 1,
            'errors': []
        },
        {
            'successful_count': 0,
            'errors': [item2]
        },
    ]
    result = script.move_valid_statement_attachments(
        env, source_dir, dest_dir, limit)
    assert result == expected_result
    get_data_from_dynamodb.assert_called_once_with(
        env, source_dir, limit)
    process_items.assert_has_calls([
        mock.call(env=env, items=expected_items, dest_dir=dest_dir),
        mock.call(env=env, items=expected_items2, dest_dir=dest_dir),
    ])


@mock.patch('move_attachments_per_label_dir.move_valid_statement_attachments')
@mock.patch('move_attachments_per_label_dir.get_estimated_time')
@mock.patch('move_attachments_per_label_dir.get_approximate_table_count')
@mock.patch('move_attachments_per_label_dir.configure_logger')
@mock.patch('move_attachments_per_label_dir.common_config')
def test_main(
        common_config, configure_logger, get_approximate_table_count,
        get_estimated_time, move_valid_statement_attachments):
    """Test script entry point main."""
    limit = 200
    dest_dir = 'expected_dest'
    source_dir = 'source_test'
    env = 'testenv'
    args = mock.MagicMock(
        log_level='info', scan_limit=limit, dest_dir=dest_dir,
        source_dir=source_dir, env=env)
    item_count = 42
    get_approximate_table_count.return_value = item_count
    get_estimated_time.return_value = '00:00:00'
    expected_move_result = {
        'ScannedCount': 10,
        'Count': 1,
        'successful_count': 1,
        'errors': []
    }
    move_valid_statement_attachments.return_value = expected_move_result

    script.main(args)
    get_approximate_table_count.assert_called_once_with(env)
    get_estimated_time.assert_called_once_with(item_count, limit)
    move_valid_statement_attachments.assert_called_once_with(
        env=env,
        source_dir=source_dir,
        dest_dir=dest_dir,
        limit=limit)