"""Unit tests for tasks of Amazon Unlimited Source of Stream Workflow."""

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

import pytest

from analytics_aggregation.flows.amazon_unlimited_sos import tasks
from analytics_aggregation.util import common


def test_extract_flow_params(monkeypatch):
    """Test extract_flow_params."""
    # Mocking
    feed_name = 'feed_name'
    source_feed_name = 'sourcre_feed_name'
    context_date_range = 'context_date_range'
    days_back = 1
    reload = True
    labelids = [1, 2]
    activity_mock = Mock()

    reload_result = True
    get_bool_from_flag_mock = Mock(return_value=reload_result)
    monkeypatch.setattr(common, 'get_bool_from_flag', get_bool_from_flag_mock)

    validate_date_range_reload_mock = Mock()
    monkeypatch.setattr(
        common, 'validate_date_range_reload', validate_date_range_reload_mock)

    start_date, end_date = 'start_date', 'end_date'
    find_flow_date_range_mock = Mock(return_value=(start_date, end_date))
    monkeypatch.setattr(
        common, 'find_flow_date_range', find_flow_date_range_mock)

    format_id_list_result = [2, 1]
    format_id_list_mock = Mock(return_value=format_id_list_result)
    monkeypatch.setattr(common, 'format_id_list', format_id_list_mock)

    expected_context = {
        'reload': reload_result,
        'date_range_as_str': start_date + '_' + end_date,
        'date_range': {
            'start_date': start_date,
            'end_date': end_date},
        'labelids': format_id_list_result}

    # Tested function call
    result_context = tasks.bootstrap(
        activity_mock, feed_name, source_feed_name, context_date_range,
        days_back, reload, labelids)

    # Checks
    assert result_context == expected_context
    get_bool_from_flag_mock.assert_called_once_with(reload)
    validate_date_range_reload_mock.assert_called_once_with(
        context_date_range, reload)
    find_flow_date_range_mock.assert_called_once_with(
        feed_name, source_feed_name, context_date_range, days_back)
    format_id_list_mock.assert_called_once_with(labelids)


def test_extract_flow_params_failure(monkeypatch):
    """Test extract_flow_params."""
    # Mocking
    feed_name = 'feed_name'
    source_feed_name = 'sourcre_feed_name'
    context_date_range = 'context_date_range'
    days_back = 1
    reload = True
    labelids = [1, 2]
    activity_mock = Mock()

    reload_result = True
    get_bool_from_flag_mock = Mock(return_value=reload_result)
    monkeypatch.setattr(common, 'get_bool_from_flag', get_bool_from_flag_mock)

    validate_date_range_reload_mock = Mock()
    error_message = 'error_message'
    validate_date_range_reload_mock.side_effect = (
        common.UnprocessableFlowParamsException(error_message))
    monkeypatch.setattr(
        common, 'validate_date_range_reload', validate_date_range_reload_mock)

    expected_message = 'expected_message'
    expected_message_mock = Mock(return_value=expected_message)
    monkeypatch.setattr(common, 'exit_message', expected_message_mock)

    # Tested function call
    result = tasks.bootstrap(
        activity_mock, feed_name, source_feed_name, context_date_range,
        days_back, reload, labelids)

    # Checks
    assert result == expected_message


@pytest.yield_fixture
def mock_executor_context():
    """Yield executor context."""
    sf_executor_class_path = (
        'analytics_aggregation.flows.amazon_unlimited_sos.'
        'tasks.AmazonUnlimitedSF')
    with patch(sf_executor_class_path) as sf_executor:
        mock_executor_context = sf_executor.return_value.__enter__.return_value
        yield mock_executor_context


def test_cleanup_staging_sos(mock_executor_context):
    """Test cleanup_staging_sos."""
    # Mocking
    start_date = 'start_date'
    end_date = 'end_date'
    date_range = {'start_date': start_date, 'end_date': end_date}
    labelids = [1, 2, 3]
    # Tested function call
    tasks.cleanup_staging_sos(Mock(), date_range, labelids)
    # Checks
    mock_executor_context.cleanup_staging_sos.assert_called_once_with(
        date_range, labelids)


def test_populate_staging_sos(mock_executor_context):
    """Test populate_staging_sos."""
    # Mocking
    start_date = 'start_date'
    end_date = 'end_date'
    date_range = {'start_date': start_date, 'end_date': end_date}
    labelids = [1, 2, 3]
    # Tested function call
    tasks.populate_staging_sos(Mock(), date_range, labelids)
    # Checks
    mock_executor_context.populate_staging_sos.assert_called_once_with(
        date_range, labelids)