"""Unit tests for snowflake_flatten_contract_json task."""

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

from lib import config
from lib import constants
from tasks.accounting_run_calculate.snowflake_flatten_contract_json \
    import flatten_contract_json_task


import_path = 'tasks.accounting_run_calculate.snowflake_flatten_contract_json'


@patch(f'{import_path}.RoyaltySnowflakeHook')
@patch(f'{import_path}.helpers')
@patch(f'{import_path}.aws')
@patch(f'{import_path}.insert_flat_distro_contract_data_template')
def test_flatten_contract_json_task(
    mock_template,
    mock_aws,
    mock_helpers,
    mock_hook,
    mock_accounting_run_calculate_dag_run
):
    """Test task that flattens contract JSON and inserts results into snowflake."""
    accounting_run_id = 123
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'Test Period'
    }
    accounting_run = {
        'accounting_period_id': 1,
        'accounting_run_id': accounting_run_id,
        'run_controller_name': 'Test Run'
    }
    base_key = '1-test-period/123-test-run'

    snapshots_url = \
        f'{base_key}' \
        f'/{constants.DIRECTORY_MODULE_MATCH_CONTRACTS}' \
        f'/{constants.DIRECTORY_SNAPSHOTS}' \

    contract_csv_url = f'{snapshots_url}/{constants.FILE_NAME_CONTRACTS_CSV}'
    term_csv_url = f'{snapshots_url}/{constants.FILE_NAME_CONTRACT_TERMS_CSV}'

    mock_helpers.get_event_from_params.return_value.target_id = accounting_run_id
    mock_helpers.get_event_records.return_value = (accounting_period, accounting_run)
    mock_helpers.build_snapshot_key_from_accounting_run.return_value = snapshots_url

    mock_aws.location.side_effect = [
        MagicMock(key=contract_csv_url),
        MagicMock(key=term_csv_url)
    ]
    mock_template.return_value.render.return_value = \
        'FLATTEN CONTRACT JSON AND INSERT INTO CONTRACT_DENORMALIZED_DISTRO'
    mock_hook.return_value.run.return_value = True

    flatten_contract_json_task(mock_accounting_run_calculate_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_accounting_run_calculate_dag_run
    )
    mock_helpers.get_event_records.assert_called_once_with(accounting_run_id)
    mock_helpers.build_snapshot_key_from_accounting_run.assert_called_once_with(
        accounting_period,
        accounting_run,
        constants.DIRECTORY_MODULE_MATCH_CONTRACTS
    )
    mock_aws.location.assert_any_call(
        snapshots_url,
        constants.FILE_NAME_CONTRACTS_CSV
    )
    mock_aws.location.assert_called_with(
        snapshots_url,
        constants.FILE_NAME_CONTRACT_TERMS_CSV
    )
    mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)
    mock_template.return_value.render.assert_called_once_with(
        accounting_run_id=accounting_run_id,
        contract_csv_url=contract_csv_url,
        schema=config.OWS_ENV,
        tsv_stage=config.ABACUS_TSV_STAGE,
        term_csv_url=term_csv_url
    )
    mock_hook.return_value.run.assert_called_once()