"""Test for invoke_reserves_schedule_lambda task."""

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

import pytest

from hooks.lambda_invoke.invocation_response import LambdaInvocationResponse
from lib.constants import CONTRACT_TYPES
from lib.constants import DAG_SCHEDULE_RESERVES
from tasks.accounting_run_commit.invoke_reserves_schedule_lambda import \
    invoke_reserves_schedule_lambda_task


@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.OrchLambdaHook')
@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.helpers')
def test_invoke_reserves_schedule_lambda_task_distribution_success(
    mock_helpers,
    mock_orch_lambda_hook,
    mock_commit_dag_run,
    lambda_response_success,
    mock_schedule_reserves_event
):
    """Test task invokes reserves_schedule lambda for distribution accounting period."""
    accounting_run_id = mock_schedule_reserves_event.get('target_id')
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'A Distribution Period',
        'contract_type': CONTRACT_TYPES.DISTRIBUTION
    }
    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)
    mock_helpers.create_event.return_value.to_json.return_value = \
        mock_schedule_reserves_event

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

    invoke_reserves_schedule_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,
        DAG_SCHEDULE_RESERVES
    )
    mock_orch_lambda_hook.assert_called_once()
    mock_orch_lambda_hook.return_value.invoke_lambda.assert_called_once_with(
        mock_schedule_reserves_event
    )


@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.OrchLambdaHook')
@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.helpers')
def test_invoke_reserves_schedule_lambda_task_nr_success(
    mock_helpers,
    mock_orch_lambda_hook,
    mock_commit_dag_run,
    mock_schedule_reserves_event
):
    """Test task does not invoke reserves_take lambda for NR accounting_period."""
    accounting_run_id = mock_schedule_reserves_event.get('target_id')
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'A Neighbouring Rights Period',
        'contract_type': 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_schedule_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_orch_lambda_hook.assert_not_called()
    mock_orch_lambda_hook.return_value.invoke_lambda.assert_not_called()


@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.OrchLambdaHook')
@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.helpers')
def test_invoke_reserves_schedule_lambda_task_error(
    mock_helpers,
    mock_orch_lambda_hook,
    mock_commit_dag_run,
    lambda_response_error,
    mock_schedule_reserves_event
):
    """Test for invoke_reserves_schedule_lambda task."""
    accounting_run_id = mock_schedule_reserves_event.get('target_id')
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'A Distribution Period',
        'contract_type': CONTRACT_TYPES.DISTRIBUTION
    }
    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)
    mock_helpers.create_event.return_value.to_json.return_value = \
        mock_schedule_reserves_event

    mock_lambda_response = LambdaInvocationResponse(lambda_response_error)
    mock_orch_lambda_hook.return_value.invoke_lambda.return_value = mock_lambda_response

    with pytest.raises(Exception):
        invoke_reserves_schedule_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,
        DAG_SCHEDULE_RESERVES
    )
    mock_orch_lambda_hook.assert_called_once()
    mock_orch_lambda_hook.return_value.invoke_lambda.assert_called_once_with(
        mock_schedule_reserves_event
    )


@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.OrchLambdaHook')
@patch('tasks.accounting_run_commit.invoke_reserves_schedule_lambda.helpers')
def test_invoke_reserves_schedule_lambda_task_get_contract_type(
    mock_helpers,
    mock_orch_lambda_hook,
    mock_commit_dag_run,
    mock_schedule_reserves_event
):
    """Test when task gets contract_type from xcom it does not use helper method."""
    accounting_run_id = mock_schedule_reserves_event.get('target_id')
    accounting_period = {
        'accounting_period_id': 1,
        'accounting_period_name': 'A Neighbouring Rights Period',
        'contract_type': 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 = 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_schedule_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_orch_lambda_hook.assert_not_called()
    mock_orch_lambda_hook.return_value.invoke_lambda.assert_not_called()