"""Unit tests for Amazon Unlimited Snowflake SQL executor."""

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

import pytest

from analytics_aggregation.flows.amazon_unlimited_sos.snowflake_executor \
    import AmazonUnlimitedSF
from analytics_aggregation.util import common


@pytest.yield_fixture
def sql_loader_mock():
    """Return sql_loader mock."""
    sql_loader_path = (
        'analytics_aggregation.flows.amazon_unlimited_sos.'
        'snowflake_executor.sql_loader')
    with patch(sql_loader_path) as sql_loader:
        yield sql_loader


@pytest.yield_fixture
def executor_mock(sf_config_mock):
    """Yield executor context."""
    executor = AmazonUnlimitedSF(sf_config_mock)
    with patch.object(executor, 'execute', wraps=executor.execute) as \
            executor.ex_mock:
        yield executor


def test_execute_dynamic_sql(sql_loader_mock, executor_mock):
    """Test _execute_dynamic_sql method."""
    # Mocking
    template_name = 'template_name'
    sql_params = {'test': 'param'}
    format_params = {'test': 'param'}
    sql_template_mock = Mock()
    sql_loader_mock.load_query.return_value = sql_template_mock
    formated_sql_template = Mock()
    sql_template_mock.format.return_value = formated_sql_template
    sql_mock = Mock()
    non_identifier_params_mock = Mock()
    validator_mock = Mock()
    executor_mock.validator = validator_mock
    executor_mock.validator.format_identifiers.return_value = (
        sql_mock, non_identifier_params_mock)
    # Tested method call
    executor_mock._execute_dynamic_sql(
        template_name, sql_params, format_params)
    # Checks
    sql_loader_mock.load_query.assert_called_once_with(template_name)
    executor_mock.validator.format_identifiers.assert_called_once_with(
        formated_sql_template, sql_params)
    executor_mock.execute.assert_called_once_with(
        sql_mock, params=non_identifier_params_mock)


def test_cleanup_staging_sos(monkeypatch, executor_mock):
    """Test cleanup_staging_sos method."""
    # Mocking
    date_range = {'start_date': 'start_date', 'end_date': 'end_dage'}
    labelids = [1, 2]
    sql_params = dict(
        schema=executor_mock.sf_config['schema'],
        start_date=date_range['start_date'],
        end_date=date_range['end_date'],
        labelids=labelids)
    sos_labelid_clause_mock = Mock()
    labelids_clause = 'labelids_clause'
    sos_labelid_clause_mock.return_value = labelids_clause
    monkeypatch.setattr(common, 'sos_labelid_clause', sos_labelid_clause_mock)
    template_format_params = dict(labelids_clause=labelids_clause)
    execute_dynamic_sql_mock = Mock()
    executor_mock._execute_dynamic_sql = execute_dynamic_sql_mock
    # Tested function call
    executor_mock.cleanup_staging_sos(date_range, labelids)
    # Checks
    sos_labelid_clause_mock.assert_called_once_with(labelids)
    executor_mock._execute_dynamic_sql.assert_called_once_with(
        'delete_from_staging_sos', sql_params, template_format_params)


def test_populate_staging_sos(monkeypatch, executor_mock):
    """Test cleanup_staging_sos method."""
    # Mocking
    date_range = {'start_date': 'start_date', 'end_date': 'end_dage'}
    labelids = [1, 2]
    sql_params = dict(
        schema=executor_mock.sf_config['schema'],
        start_date=date_range['start_date'],
        end_date=date_range['end_date'],
        labelids=labelids)
    sos_labelid_clause_mock = Mock()
    labelids_clause = 'labelids_clause'
    sos_labelid_clause_mock.return_value = labelids_clause
    monkeypatch.setattr(common, 'sos_labelid_clause', sos_labelid_clause_mock)
    template_format_params = dict(labelids_clause=labelids_clause)
    execute_dynamic_sql_mock = Mock()
    executor_mock._execute_dynamic_sql = execute_dynamic_sql_mock
    # Tested function call
    executor_mock.populate_staging_sos(date_range, labelids)
    # Checks
    sos_labelid_clause_mock.assert_called_once_with(labelids, 'dt')
    executor_mock._execute_dynamic_sql.assert_called_once_with(
        'populate_staging_with_amazon_unlimited_data', sql_params,
        template_format_params)