"""Unit tests for tasks of YouTube Claim Ingestion Workflow."""
import datetime
from unittest.mock import MagicMock
from unittest.mock import patch

from freezegun import freeze_time
from garcon_contrib.dynamo_feed_status import garcon_feed_status
import pytest

from feed_ingestion.flows.youtube_claim import config
from feed_ingestion.flows.youtube_claim import tasks


@pytest.fixture(params=[
    'theorchard'
])
def context_and_params(request):
    """Yield context, licensor."""
    licensor = request.param
    context = {
        'activity': MagicMock(),
        'date': '2019-11-01',
        'licensor': licensor,
        'reload': False}
    yield context, licensor


@pytest.fixture
def expected_bootstrap_response(context_and_params):
    """Response for bootstrap task."""
    context, licensor = context_and_params
    feed_name = config.feed_name + '_{licensor}'.format(licensor=licensor)
    if licensor == 'theorchard':
        temp_table_names = [
            config.temp_table_name_template.format(
                staging_raw_table=(
                    config.snowflake_table_names['staging_raw_theorchard']),
                date=datetime.date(2018, 3, 31),
                account=account
            ) for account in config.cms_dict.values()
        ]
        cms_dict = config.cms_dict

    yield dict(
        feed_name=feed_name,
        licensor=licensor,
        report_name=config.youtube_report_full_name,
        secrets_path=config.secrets_path,
        date='2018-03-31',
        report_date='2018-03-31',
        processed_datetime='2018-03-31T00:00:00',
        s3_dir_path='Youtube_claims/archives/2018-03-31/{licensor}/'.format(
            licensor=licensor),
        s3_archive_path='s3://dev-cucumbers/'
                        'Youtube_claims/archives/'
                        '2018-03-31/{licensor}/'.format(licensor=licensor),
        temp_table_names=temp_table_names,
        s3_split_path=('s3://dev-cucumbers/'
                       'temp-feed-ingestion/Youtube_claims/'
                       '2018-03-31/theorchard/'),
        credentials_path=config.credentials_paths[licensor],
        api_service_name='youtubereporting',
        api_version='v1',
        jobs_meta_path=config.jobs_meta_path,
        cms_dict=cms_dict,
        staging_raw_table_name='staging_raw_youtube_claim',
        load_temp_staging_raw_table_sme=(
            'staging_raw_youtube_claim_20180331_sme'),
    )


@pytest.fixture
def mock_task_status():
    """Yield task status."""
    task_status_path = 'feed_ingestion.flows.youtube_claim.tasks.task_status'
    with patch(task_status_path) as task_status:
        task_status.is_completed_task.return_value = False
        task_status.mark_completed_task = MagicMock()
        yield task_status


@pytest.fixture
def mock_check_status():
    """Yield task status."""
    task_status_path = 'feed_ingestion.tasks.task_status'
    with patch(task_status_path) as task_status:
        task_status.is_completed_task.return_value = False
        task_status.mark_completed_task = MagicMock()
        yield task_status


@pytest.fixture
def mock_set_overall_status():
    """Yield overall status."""
    overall_status_path = (
        'feed_ingestion.flows.youtube_claim.tasks.garcon_feed_status.'
        'set_overall_status')
    with patch(overall_status_path) as overall_status:
        yield overall_status


@pytest.fixture
def mock_executor_context():
    """Yield executor context."""
    reg_executors_class_path = (
        'feed_ingestion.flows.youtube_claim.tasks.registered_executors')
    with patch(reg_executors_class_path) as reg_executors:
        mock_executor = MagicMock()
        reg_executors.get = mock_executor
        yield mock_executor.return_value.return_value.__enter__.return_value


@freeze_time('2018-03-31')
@patch('feed_ingestion.flows.youtube_claim.'
       'tasks.garcon_feed_status.get_overall_status')
def test_bootstrap(mock_get_overall_status, expected_bootstrap_response):
    """Check that bootstrap returns expected results."""
    licensor = expected_bootstrap_response['licensor']
    context = {
        'activity': MagicMock(),
        'date': '2018-03-31',
        'reload': True,
        'licensor': licensor
    }
    reply = tasks.bootstrap(**context)
    feed_name = expected_bootstrap_response['feed_name']

    mock_get_overall_status.assert_called_with(
        feed_name, '2018-03-31')
    assert reply == expected_bootstrap_response


def test_clean_staging_raw_table(
        mock_check_status, mock_executor_context, expected_bootstrap_response):
    """Test clean_staging_raw_table task."""
    feed_name = expected_bootstrap_response['feed_name']
    mock_check_status.return_value = False
    tasks.clean_staging_raw_table(
        MagicMock(), '2018-03-31', 'staging_raw_table', feed_name)
    mock_executor_context.clean_staging_raw_table.assert_called_with(
        'staging_raw_table', '2018-03-31')


def test_create_temp_staging_raw_table(
        mock_task_status, mock_executor_context, expected_bootstrap_response):
    """Test create_temp_staging_raw_table task."""
    call_count = len(expected_bootstrap_response['temp_table_names'])
    tasks.create_temp_staging_raw_table(
        MagicMock(), '2018-03-31',
        expected_bootstrap_response['temp_table_names'],
        expected_bootstrap_response['feed_name'])
    assert mock_task_status.is_completed_task.call_count == call_count
    assert mock_task_status.mark_completed_task.call_count == call_count
    assert mock_executor_context.create_temp_staging_raw_table.call_count \
        == call_count


def test_load_temp_staging_raw_table(
        mock_task_status, mock_executor_context, expected_bootstrap_response):
    """Test load_temp_staging_raw_table task."""
    call_count = len(expected_bootstrap_response['temp_table_names'])
    tasks.load_temp_staging_raw_table(
        MagicMock(), '2018-03-31',
        expected_bootstrap_response['s3_archive_path'],
        expected_bootstrap_response['temp_table_names'],
        expected_bootstrap_response['feed_name'])
    assert mock_task_status.is_completed_task.call_count == call_count
    assert mock_task_status.mark_completed_task.call_count == call_count
    assert mock_executor_context.load_temp_staging_raw_table.call_count \
           == call_count


@patch('feed_ingestion.flows.youtube_claim.tasks.s3.get_key_size')
def test_load_staging_raw_table(
        mock_get_key_size, mock_task_status, mock_executor_context,
        expected_bootstrap_response):
    """Test load_staging_raw_table task."""
    call_count = len(expected_bootstrap_response['temp_table_names'])
    tasks.load_staging_raw_table(
        MagicMock(), '2018-03-31',
        expected_bootstrap_response['processed_datetime'],
        expected_bootstrap_response['temp_table_names'],
        'staging_raw_table',
        expected_bootstrap_response['s3_archive_path'],
        expected_bootstrap_response['feed_name'])
    assert mock_task_status.is_completed_task.call_count == 1
    assert mock_task_status.mark_completed_task.call_count == 2
    assert mock_executor_context.load_staging_raw_table.call_count \
        == call_count


def test_drop_temp_table(
        mock_task_status, mock_executor_context, expected_bootstrap_response):
    """Test drop_temp_table task."""
    call_count = len(expected_bootstrap_response['temp_table_names'])
    tasks.drop_temp_table(
        MagicMock(),
        expected_bootstrap_response['temp_table_names'],
        expected_bootstrap_response['feed_name'])
    assert mock_executor_context.drop_table.call_count \
        == call_count


def test_check_assets_availability(
        mock_set_overall_status, mock_executor_context,
        expected_bootstrap_response):
    """Test check_assets_availability task."""
    feed_name = expected_bootstrap_response['feed_name']
    mock_executor_context.is_assets_available.return_value = 1
    tasks.check_assets_availability(MagicMock(), '2018-06-13', feed_name)
    assert not mock_set_overall_status.called

    mock_executor_context.is_assets_available.return_value = 0
    tasks.check_assets_availability(MagicMock(), '2018-06-13', feed_name)
    mock_set_overall_status.assert_called_with(
        feed_name, '2018-06-13',
        garcon_feed_status.STATUS_NOT_AVAILABLE)


def test_update_dim_claim(
        mock_task_status, mock_executor_context, expected_bootstrap_response):
    """Test update_claim task."""
    tasks.update_dim_claim(MagicMock(),
                           '2020-01-01',
                           expected_bootstrap_response['feed_name'])
    assert mock_executor_context.update_dim_claim.call_count == 1
    assert mock_task_status.mark_completed_task.call_count == 1