"""Unit tests for updating the `adjustment_file_generate` Abacus states."""

from unittest.mock import patch

from lib.constants import (
    ABACUS_STATE_STATUSES,
    AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS,
)
from tasks.adjustment_file_generate import states


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_start_generate_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "generate" action as running."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.INIT
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    states.start_generate_task(mock_adjustment_file_generate_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.RUNNING
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_success_generate_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "generate" action as complete."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.RUNNING
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    states.success_generate_task(mock_adjustment_file_generate_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.COMPLETE
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_failure_generate_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "generate" action as failed."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.RUNNING
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    result = states.failure_generate_task(mock_adjustment_file_generate_dag_run)

    assert result is False

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.ERROR
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_start_validate_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "validate" action as running."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.INIT
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    states.start_validate_task(mock_adjustment_file_generate_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.RUNNING
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_success_validate_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "validate" action as complete."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.RUNNING
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    states.success_validate_task(mock_adjustment_file_generate_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.COMPLETE
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_failure_validate_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "validate" action as failed."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.RUNNING
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    result = states.failure_validate_task(mock_adjustment_file_generate_dag_run)

    assert result is False

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.ERROR
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_start_import_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "import" action as running."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.INIT
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    states.start_import_task(mock_adjustment_file_generate_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.RUNNING
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_success_import_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "import" action as complete."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.RUNNING
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    states.success_import_task(mock_adjustment_file_generate_dag_run)

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.COMPLETE
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )


@patch('tasks.adjustment_file_generate.states.helpers')
@patch('tasks.adjustment_file_generate.states.ows')
def test_failure_import_task(
    mock_ows,
    mock_helpers,
    mock_adjustment_file_generate_dag_run
):
    """Test marking the "import" action as failed."""
    target_id = 1

    mock_action = {
        'abacus_state_id': 1,
        'action_name': (
            AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE
        ),
        'action_status': ABACUS_STATE_STATUSES.RUNNING
    }

    mock_helpers.get_event_from_params.return_value.target_id = target_id
    mock_helpers.get_abacus_state.return_value = mock_action
    mock_ows.update_abacus_state.return_value = True

    result = states.failure_import_task(mock_adjustment_file_generate_dag_run)

    assert result is False

    mock_helpers.get_event_from_params.assert_called_once_with(
        mock_adjustment_file_generate_dag_run
    )

    mock_helpers.get_abacus_state.assert_called_once_with(
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE,
        target_id
    )

    mock_ows.update_abacus_state.assert_called_once_with(
        1,
        body={
            'action_status': ABACUS_STATE_STATUSES.ERROR
        },
        headers=states.DEFAULT_UPDATE_HEADERS
    )