"""Test snowflake_export_sales task."""

from unittest.mock import MagicMock
from unittest.mock import patch

from lib import config
from tasks.sales_approve.snowflake_export_sales import export_sales_from_snowflake


@patch('tasks.sales_approve.snowflake_export_sales.RoyaltySnowflakeHook')
@patch('tasks.sales_approve.snowflake_export_sales.helpers')
@patch('tasks.sales_approve.snowflake_export_sales.export_approved_distro_sales_to_s3')
@patch('tasks.sales_approve.snowflake_export_sales.ows')
def test_export_distribution_sales_from_snowflake(
    mock_ows,
    mock_template,
    mock_helpers,
    mock_hook,
    mock_sales_approve_event,
    mock_sales_approve_dag_run
):
    """Test building s3 path, exporting distribution sales from SF to s3, and updating sales file."""  # noqa: E501
    sales_file_id = mock_sales_approve_event.get('target_id')

    accounting_period = {
        'accounting_period_id': 123,
        'accounting_period_name': 'Month Year',
        'contract_type': 'distribution'
    }
    sales_file = {
        'sales_file_id': sales_file_id,
        'file_name': 'make it rain'
    }

    s3_key = f'123-month-year/eligible-sales/{sales_file_id}-make-it-rain/results/'
    s3_url = f's3://{s3_key}'

    mock_event = MagicMock(target_id=sales_file_id)
    mock_s3_location = MagicMock(key=s3_key, url=s3_url)

    mock_helpers.get_event_from_params.return_value = mock_event
    mock_helpers.get_event_records.return_value = (sales_file, accounting_period)
    mock_helpers.build_sales_file_main_url.return_value = mock_s3_location

    mock_template.return_value.render.return_value = 'COPY SALES TO S3 AS PARQUET'
    mock_hook.return_value.run.return_value = True

    mock_ows.update_sales_file.return_value = {'main_url': s3_url, **sales_file}

    export_sales_from_snowflake(mock_sales_approve_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_sales_approve_dag_run
    )
    mock_helpers.get_event_records.assert_called_once_with(mock_event)
    mock_helpers.build_sales_file_main_url.assert_called_once_with(
        accounting_period,
        sales_file
    )
    mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)
    mock_template.return_value.render.assert_called_once_with(
        env=config.OWS_ENV,
        s3_path=s3_key,
        sales_file_id=sales_file_id,
        stage=config.ABACUS_PARQUET_STAGE
    )
    mock_hook.return_value.run.assert_called_once()
    mock_ows.update_sales_file.assert_called_once_with(
        sales_file_id,
        **{'main_url': s3_url}
    )


@patch('tasks.sales_approve.snowflake_export_sales.RoyaltySnowflakeHook')
@patch('tasks.sales_approve.snowflake_export_sales.helpers')
@patch('tasks.sales_approve.snowflake_export_sales.export_approved_nr_sales_to_s3')
@patch('tasks.sales_approve.snowflake_export_sales.ows')
def test_export_nr_sales_from_snowflake(
    mock_ows,
    mock_nr_sales_template,
    mock_helpers,
    mock_hook,
    mock_sales_approve_event,
    mock_sales_approve_dag_run
):
    """Test building s3 path, exporting NR sales from SF to s3, and updating sales file."""  # noqa: E501
    sales_file_id = mock_sales_approve_event.get('target_id')

    accounting_period = {
        'accounting_period_id': 123,
        'accounting_period_name': 'Month Year',
        'contract_type': 'neighbouring_rights'
    }
    sales_file = {
        'sales_file_id': sales_file_id,
        'file_name': 'make it rain'
    }

    s3_key = f'123-month-year/eligible-sales/{sales_file_id}-make-it-rain/results/'
    s3_url = f's3://{s3_key}'

    mock_event = MagicMock(target_id=sales_file_id)
    mock_s3_location = MagicMock(key=s3_key, url=s3_url)

    mock_helpers.get_event_from_params.return_value = mock_event
    mock_helpers.get_event_records.return_value = (sales_file, accounting_period)
    mock_helpers.build_sales_file_main_url.return_value = mock_s3_location

    mock_nr_sales_template.return_value.render.return_value = 'COPY SALES TO S3 AS PARQUET'  # noqa: E501
    mock_hook.return_value.run.return_value = True

    mock_ows.update_sales_file.return_value = {'main_url': s3_url, **sales_file}

    export_sales_from_snowflake(mock_sales_approve_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_sales_approve_dag_run
    )
    mock_helpers.get_event_records.assert_called_once_with(mock_event)
    mock_helpers.build_sales_file_main_url.assert_called_once_with(
        accounting_period,
        sales_file
    )
    mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)
    mock_nr_sales_template.return_value.render.assert_called_once_with(
        env=config.OWS_ENV,
        s3_path=s3_key,
        sales_file_id=sales_file_id,
        stage=config.ABACUS_PARQUET_STAGE
    )
    mock_hook.return_value.run.assert_called_once()
    mock_ows.update_sales_file.assert_called_once_with(
        sales_file_id,
        **{'main_url': s3_url}
    )