"""Tests for producer script."""

from datetime import datetime

import config
import const
import producer

import boto3
import flexmock
from moto import mock_dynamodb2
import pytest


def test_send_record():
    """Test sending an aggregated record to kinesis."""
    pk = 'ISRC12345'
    ehk = 'ehk'
    data = {'isrc': 'ISRC12345', 'territories': {'US': 12344}}
    kinesis_client = flexmock()
    kinesis_client.should_receive('put_record').once()
    (flexmock(producer).
        should_receive('kinesis_client').
        and_return(kinesis_client))
    agg_record = flexmock(get_num_user_records=lambda: 100)
    (agg_record.
        should_receive('get_contents').
        and_return((pk, ehk, data)).
        once())
    producer.send_record(agg_record)


@pytest.mark.parametrize('records, on_time', [
    ((1), True), ((0, 1), False)
])
def test_wait_for_finish(records, on_time):
    """Test wait for finish of the export job."""
    job_id = 'some_id2233'
    redis_mock = flexmock(hdel=lambda a, b: True)
    redis_mock.should_receive('hget').and_return(records).one_by_one()
    flexmock(config).should_receive('TIMEOUT').and_return(1)
    (flexmock(producer).
        should_receive('decrease_read_capacity').
        and_return(True))
    assert (
        producer.wait_for_finish(
            job_id, redis_mock, datetime.utcnow(), 1) == on_time)


def test_scan_dynamo(mocker):
    """Test scanning the dynamoDB table."""
    mock_dynamo_client = flexmock()
    mock_dynamo_result = {
        'Items': ['isrc1', 'isrc2'],
        'LastEvaluatedKey': 'key'
    }
    mock_dynamo_result2 = {
        'Items': ['isrc3', 'isrc4'],
        'LastEvaluatedKey': None
    }
    (mock_dynamo_client.
        should_receive('scan').
        and_return(mock_dynamo_result, mock_dynamo_result2).
        one_by_one())

    flexmock(producer).should_receive('send_to_kinesis').with_args(
        list, 'abc1').twice()

    kinesis_agg = flexmock(clear_and_get=lambda: True)
    flexmock(producer).should_receive('send_record').and_return(True)

    mock_lock = flexmock()
    flexmock(producer).should_receive('aggregator_lock').and_return(mock_lock)

    mock_lock.should_receive('acquire').at_least().once()
    mock_lock.should_receive('release').at_least().once()

    producer.kinesis_agg = kinesis_agg
    producer.scan_dynamo(0, 1, mock_dynamo_client, 'abc1')


READ_CAPACITY = 20
WRITE_CAPACITY = 5


@mock_dynamodb2
def patch_dynamo_active_table(read=READ_CAPACITY, write=WRITE_CAPACITY):
    """Create mock dynamo DB table."""
    dynamo_resource = boto3.resource('dynamodb', 'us-east-1')
    dynamo_resource.create_table(
        AttributeDefinitions=[
            {
                'AttributeName': 'isrc',
                'AttributeType': 'S'
            }
        ],
        TableName='test-masters_active',
        KeySchema=[
            {
                'AttributeName': 'isrc',
                'KeyType': 'HASH'
            }
        ],
        ProvisionedThroughput={
            'ReadCapacityUnits': read,
            'WriteCapacityUnits': write
        }
    )
    mock_table = dynamo_resource.Table('test-masters_active')
    producer.active_table = mock_table
    return mock_table


@mock_dynamodb2
def test_change_dynamo_read_cap():
    """Assert that dynamoDB is called with the correct arguments."""
    read_cap = 10

    active_table = patch_dynamo_active_table()

    (flexmock(active_table).should_call('update').with_args(
        ProvisionedThroughput={
            const.READ_CAPACITY_UNITS: read_cap,
            const.WRITE_CAPACITY_UNITS: WRITE_CAPACITY
        }
    ).once())
    producer.change_dynamo_read_cap(read_cap)


@mock_dynamodb2
def test_increase_read_capacity():
    """Test increase DynamoDB read capacity."""
    read_cap = READ_CAPACITY + config.DDB_ADDITIONAL_READ_CAP

    active_table = patch_dynamo_active_table()

    (flexmock(active_table).should_call('update').with_args(
        ProvisionedThroughput={
            const.READ_CAPACITY_UNITS: read_cap,
            const.WRITE_CAPACITY_UNITS: WRITE_CAPACITY
        }
    ).once())
    producer.increase_read_capacity(producer.active_table)


@mock_dynamodb2
def test_decrease_read_capacity():
    """Test decrease DynamoDB read capacity."""
    active_table = patch_dynamo_active_table(
        READ_CAPACITY + config.DDB_ADDITIONAL_READ_CAP)

    (flexmock(active_table).should_call('update').with_args(
        ProvisionedThroughput={
            const.READ_CAPACITY_UNITS: READ_CAPACITY,
            const.WRITE_CAPACITY_UNITS: WRITE_CAPACITY
        }
    ).once())
    producer.decrease_read_capacity(producer.active_table)


def test_send_to_kinesis(active_record_fixture):
    """Test send aggregated record to Kinesis stream."""
    job_id = 'abc1'
    items = [active_record_fixture]
    producer.rows_processed = 100
    aggregator_mock = flexmock()
    aggregator_mock.should_receive('acquire').once()
    aggregator_mock.should_receive('release').once()
    (flexmock(producer).
        should_receive('aggregator_lock').
        and_return(aggregator_mock))
    kinesis_agg = flexmock(clear_and_get=lambda: True)
    record = {'aggregated': 'record'}
    (kinesis_agg.
        should_receive('add_user_record').
        with_args(long, str).
        and_return(record))
    (flexmock(producer).
        should_receive('send_record').
        with_args(record).
        and_return(True))

    json_record_fixture = producer.util.dynamo_active_record_to_json(
        active_record_fixture, include_locked=True)
    (flexmock(producer.util).
        should_receive('dynamo_active_record_to_json').
        with_args(active_record_fixture, include_locked=True).
        and_return(json_record_fixture))

    producer.kinesis_agg = kinesis_agg
    producer.send_to_kinesis(items, job_id)


def test_do_export():
    """Test multi-threaded DynamoDB export."""
    total_segments = 3
    job_id = 'abc123'
    dynamo_client = flexmock()
    flexmock(producer).should_receive('scan_dynamo').and_return(True)
    log_mock = flexmock()
    flexmock(producer).should_receive('log').and_return(log_mock)
    for i in range(total_segments):
        log_mock.should_receive('info').with_args(
            'Worker {} finished.'.format(i)).at_least().once()
    (log_mock.
        should_receive('info').
        with_args('All workers finished.').
        at_least().
        once())
    producer.do_export(total_segments, dynamo_client, job_id)


def test_do_export_spawns_correct_number_of_threads_with_correct_args(mocker):
    """Expect to spawn 3 threads with correct arguments."""
    total_segments = 3
    dynamo_client = mocker.Mock()
    job_id = 'abc123'
    mock_thread_object = mocker.Mock()
    mock_thread_class = mocker.patch.object(
        producer.threading, 'Thread', return_value=mock_thread_object)
    expected_calls = [
        mocker.call(target=producer.scan_dynamo, name='Worker 0',
                    args=(0, 3, dynamo_client, job_id)),
        mocker.call(target=producer.scan_dynamo, name='Worker 1',
                    args=(1, 3, dynamo_client, job_id)),
        mocker.call(target=producer.scan_dynamo, name='Worker 2',
                    args=(2, 3, dynamo_client, job_id))]

    producer.do_export(total_segments, dynamo_client, job_id)

    assert mock_thread_class.call_count == total_segments
    mock_thread_class.assert_has_calls(expected_calls)


def test_do_export_threads_started_and_joined(mocker):
    """Expect spawned threads to be started and then joined."""
    total_segments = 3
    dynamo_client = mocker.Mock()
    job_id = 'abc123'
    mock_thread_object = mocker.Mock()
    mock_thread_class = mocker.patch.object(
        producer.threading, 'Thread', return_value=mock_thread_object)

    producer.do_export(total_segments, dynamo_client, job_id)

    assert mock_thread_class.call_count == total_segments
    assert mock_thread_object.start.call_count == total_segments
    assert mock_thread_object.join.call_count == total_segments