"""Test snowflake_check_ingest_sales task."""

from unittest.mock import MagicMock, patch

from lib import config
from tasks.sales_ingest.snowflake_check_ingested_sales import check_ingested_sales_task

import_path = 'tasks.sales_ingest.snowflake_check_ingested_sales'


@patch(f'tasks.sales_ingest.snowflake_check_ingested_sales.RoyaltySnowflakeHook')
@patch(f'tasks.sales_ingest.snowflake_check_ingested_sales.check_ingested_sales')
@patch(f'tasks.sales_ingest.snowflake_check_ingested_sales.get_abacus_event')
def test_check_ingested_sales(
    mock_get_abacus_event,
    mock_check_ingest_sales_template,
    mock_hook,
):
    """Test using snowflake hook to check ingested sales."""
    dag_run = MagicMock()
    task_instance = MagicMock()

    mock_get_abacus_event.return_value.event_name = 'ingest_sales_nr'
    mock_get_abacus_event.return_value.abacus_event_id = 123
    mock_get_abacus_event.return_value.rolled_back_at = None
    mock_get_abacus_event.return_value.target_id = 456

    task_instance.xcom_pull.return_value = 'nr'

    mock_check_ingest_sales_template.return_value.render.return_value = 'SELECT'

    mock_hook_instance = mock_hook.return_value
    sqlalchemy_engine = mock_hook_instance.get_sqlalchemy_engine.return_value
    mock_connection = sqlalchemy_engine.begin.return_value
    mock_result = mock_connection.__enter__.return_value.execute.return_value
    mock_result.fetchone.return_value = {'count': 0}

    kwargs = {'task_instance': task_instance}

    result = check_ingested_sales_task(dag_run, **kwargs)

    mock_get_abacus_event.assert_called_once_with(dag_run, **kwargs)
    mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)
    mock_check_ingest_sales_template.return_value.render.assert_called_once_with(
        env=config.OWS_ENV,
        sales_ingest_type='nr',
        batch_id=456
    )
    assert result is True
    task_instance.xcom_pull.assert_called_once_with(
        task_ids='determine_sales_ingest_type',
        key='sales_ingest_type'
    )


@patch(f'{import_path}.RoyaltySnowflakeHook')
@patch(f'{import_path}.check_ingested_sales')
@patch(f'{import_path}.get_abacus_event')
def test_check_ingested_sales_with_data(
    mock_get_abacus_event,
    mock_check_ingest_sales_template,
    mock_hook,
):
    """Test using snowflake hook to check ingested sales with data."""
    dag_run = MagicMock()
    task_instance = MagicMock()

    mock_get_abacus_event.return_value.event_name = 'ingest_sales_distro'
    mock_get_abacus_event.return_value.abacus_event_id = 123
    mock_get_abacus_event.return_value.rolled_back_at = None
    mock_get_abacus_event.return_value.target_id = 456

    task_instance.xcom_pull.return_value = 'distro'

    mock_check_ingest_sales_template.return_value.render.return_value = 'SELECT'

    mock_hook_instance = mock_hook.return_value
    sqlalchemy_engine = mock_hook_instance.get_sqlalchemy_engine.return_value
    mock_connection = sqlalchemy_engine.begin.return_value
    mock_result = mock_connection.__enter__.return_value.execute.return_value
    mock_result.fetchone.return_value = {'count': 1}

    kwargs = {'task_instance': task_instance}

    result = check_ingested_sales_task(dag_run, **kwargs)

    mock_get_abacus_event.assert_called_once_with(dag_run, **kwargs)
    mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)
    mock_check_ingest_sales_template.return_value.render.assert_called_once_with(
        env=config.OWS_ENV,
        sales_ingest_type='distro',
        batch_id=456
    )
    assert result is False