"""Test for invoke_reserves_take_lambda task."""

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

from hooks.lambda_invoke.invocation_response import LambdaInvocationResponse
from lib import constants
from tasks.accounting_run_commit.invoke_reserves_take_lambda import \
    invoke_reserves_take_lambda_task


@patch('tasks.accounting_run_commit.invoke_reserves_take_lambda.OrchLambdaHook')
@patch('tasks.accounting_run_commit.invoke_reserves_take_lambda.helpers')
def test_invoke_reserves_take_lambda_task_distribution_accounting_period(
    mock_helpers,
    mock_lambda_hook,
    mock_commit_dag_run,
    lambda_response_success
):
    """Test task invokes reserves_take lambda for distribution accounting_period."""
    accounting_run_id = 10
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'A Distribution Period',
        'contract_type': constants.CONTRACT_TYPES.DISTRIBUTION
    }
    accounting_run = {
        'accounting_run_id': accounting_run_id,
        'accounting_period_id': 1
    }
    mock_take_reserves_event = {
        'abacus_event_id': 2,
        'event_date': '2022-08-03 18:57:21.590679',
        'event_name': constants.DAG_COMMIT_TAKE_RESERVES_EVENT_NAME,
        'statement_period_id': 200,
        'target_id': accounting_run_id,
        'target_type': constants.DAG_ACCOUNTING_RUN_COMMIT_TARGET_TYPE,
        'user_id': '1',
        'user_type': 'orchard-suite'
    }

    mock_task_instance = MagicMock()
    mock_task_instance.xcom_pull.return_value = None
    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.create_event.return_value.to_json.return_value = \
        mock_take_reserves_event

    mock_lambda_response = LambdaInvocationResponse(lambda_response_success)
    mock_lambda_hook.return_value.invoke_lambda.return_value = mock_lambda_response

    invoke_reserves_take_lambda_task(
        mock_commit_dag_run,
        task_instance=mock_task_instance
    )

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_commit_dag_run,
        task_instance=mock_task_instance
    )
    mock_task_instance.xcom_pull.assert_called_once_with(
        task_ids='contract_type_branch',
        key='contract_type'
    )
    mock_helpers.get_event_records.assert_called_once_with(accounting_run_id)
    mock_helpers.create_event.assert_called_once_with(
        accounting_run_id, constants.DAG_COMMIT_TAKE_RESERVES_EVENT_NAME)
    mock_lambda_hook.assert_called_once()
    mock_lambda_hook.return_value.invoke_lambda.assert_called_once_with(
        mock_take_reserves_event
    )


@patch('tasks.accounting_run_commit.invoke_reserves_take_lambda.OrchLambdaHook')
@patch('tasks.accounting_run_commit.invoke_reserves_take_lambda.helpers')
def test_invoke_reserves_take_lambda_task_nr_accounting_period(
    mock_helpers,
    mock_lambda_hook,
    mock_commit_dag_run,
):
    """Test task does not invoke reserves_take lambda for NR accounting_period."""
    accounting_run_id = 10
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'A Neighbouring Rights Period',
        'contract_type': constants.CONTRACT_TYPES.NEIGHBOURING_RIGHTS
    }
    accounting_run = {
        'accounting_run_id': accounting_run_id,
        'accounting_period_id': 1
    }

    mock_task_instance = MagicMock()
    mock_task_instance.xcom_pull.return_value = None
    mock_helpers.get_event_from_params.return_value.target_id = accounting_run_id
    mock_helpers.get_event_records.return_value = (accounting_period, accounting_run)

    invoke_reserves_take_lambda_task(
        mock_commit_dag_run,
        task_instance=mock_task_instance
    )

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_commit_dag_run,
        task_instance=mock_task_instance
    )
    mock_task_instance.xcom_pull.assert_called_once_with(
        task_ids='contract_type_branch',
        key='contract_type'
    )
    mock_helpers.get_event_records.assert_called_once_with(accounting_run_id)
    mock_helpers.create_event.assert_not_called()
    mock_lambda_hook.assert_not_called()
    mock_lambda_hook.return_value.invoke_lambda.assert_not_called()


@patch('tasks.accounting_run_commit.invoke_reserves_take_lambda.OrchLambdaHook')
@patch('tasks.accounting_run_commit.invoke_reserves_take_lambda.helpers')
def test_invoke_reserves_take_lambda_task_get_contract_type(
    mock_helpers,
    mock_lambda_hook,
    mock_commit_dag_run,
):
    """Test when task gets contract_type from xcom it does not use helper method."""
    accounting_run_id = 10
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'A Neighbouring Rights Period',
        'contract_type': constants.CONTRACT_TYPES.NEIGHBOURING_RIGHTS
    }
    accounting_run = {
        'accounting_run_id': accounting_run_id,
        'accounting_period_id': 1
    }

    mock_task_instance = MagicMock()
    mock_task_instance.xcom_pull.return_value = \
        constants.CONTRACT_TYPES.NEIGHBOURING_RIGHTS
    mock_helpers.get_event_from_params.return_value.target_id = accounting_run_id
    mock_helpers.get_event_records.return_value = (accounting_period, accounting_run)

    invoke_reserves_take_lambda_task(
        mock_commit_dag_run,
        task_instance=mock_task_instance
    )

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_commit_dag_run,
        task_instance=mock_task_instance
    )
    mock_task_instance.xcom_pull.assert_called_once_with(
        task_ids='contract_type_branch',
        key='contract_type'
    )
    mock_helpers.get_event_records.assert_not_called()
    mock_helpers.create_event.assert_not_called()
    mock_lambda_hook.assert_not_called()
    mock_lambda_hook.return_value.invoke_lambda.assert_not_called()