from unittest import mock

import pytest

from aws_testing_utils import step_function_handler
from tests.conftest import mock_session_setup


@mock.patch('boto3.session.Session')
def test_set_machine_arn(mock_session_class: mock.MagicMock) -> None:
    name = 'testFunctionName'
    arn = 'testArn'
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_page = mock.Mock()
    mock_page.get.return_value = [{'name': name, 'stateMachineArn': arn}]
    mock_client.get_paginator.return_value.paginate.return_value = [mock_page]
    handler = step_function_handler.StepFunctionHandler(name)
    assert handler.arn == arn
    assert handler.function_name == name


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_execute_no_assert(
    mock_session_class: mock.MagicMock, mock_set_arn: mock.MagicMock
) -> None:
    mock_set_arn.return_value = 'testArn'
    response = {'executionArn': 'testExecution'}
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.start_execution.return_value = response
    handler = step_function_handler.StepFunctionHandler('test')
    handler.execute({}, False)


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.assert_execution_success'
)
@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_execute_assert(
    mock_session_class: mock.MagicMock,
    mock_set_arn: mock.MagicMock,
    mock_assert_execution_success: mock.MagicMock,
) -> None:
    mock_set_arn.return_value = 'testArn'
    response = {'executionArn': 'testExecution'}
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.start_execution.return_value = response
    mock_assert_execution_success.return_value = None
    handler = step_function_handler.StepFunctionHandler('test')
    handler.execute({}, True)


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_get_last_execution(
    mock_session_class: mock.MagicMock, mock_set_arn: mock.MagicMock
) -> None:
    arn = 'lastArn'
    mock_set_arn.return_value = 'testArn'
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.list_executions.return_value = {'executions': [{'executionArn': arn}]}
    handler = step_function_handler.StepFunctionHandler('test')
    assert handler.get_last_execution() == arn


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.get_last_execution'
)
@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_set_last_execution(
    mock_session_class: mock.MagicMock,
    mock_set_arn: mock.MagicMock,
    mock_get_last_execution: mock.MagicMock,
) -> None:
    arn = 'lastArn'
    mock_set_arn.return_value = 'testArn'
    mock_get_last_execution.return_value = arn
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    handler = step_function_handler.StepFunctionHandler('test')
    handler.set_last_execution()
    assert handler.last_execution == arn


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_assert_execution_success(
    mock_session_class: mock.MagicMock, mock_set_arn: mock.MagicMock
) -> None:
    mock_set_arn.return_value = 'testArn'
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.describe_execution.return_value = {'status': 'SUCCEEDED'}
    handler = step_function_handler.StepFunctionHandler('test')
    handler.assert_execution_success('test', 1)


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_assert_execution_error(
    mock_session_class: mock.MagicMock, mock_set_arn: mock.MagicMock
) -> None:
    mock_set_arn.return_value = 'testArn'
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.describe_execution.return_value = {'status': 'FAILED'}
    handler = step_function_handler.StepFunctionHandler('test')
    with pytest.raises(AssertionError):
        handler.assert_execution_success('test', 1)


@mock.patch('time.sleep')
@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_assert_execution_timeout(
    mock_session_class: mock.MagicMock,
    mock_set_arn: mock.MagicMock,
    mock_sleep: mock.MagicMock,
) -> None:
    mock_set_arn.return_value = 'testArn'
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.describe_execution.return_value = {'status': 'RUNNING'}
    handler = step_function_handler.StepFunctionHandler('test')
    with pytest.raises(TimeoutError):
        handler.assert_execution_success('test', 1)


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.get_last_execution'
)
@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
def test_assert_machine_ran(
    mock_set_arn: mock.MagicMock, mock_get_last_execution: mock.MagicMock
) -> None:
    mock_set_arn.return_value = 'testArn'
    arn = 'newArn'
    mock_get_last_execution.return_value = arn
    handler = step_function_handler.StepFunctionHandler('test')
    handler.last_execution = 'oldArn'
    assert handler.assert_machine_ran()


@mock.patch('time.sleep')
@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.get_last_execution'
)
@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
def test_assert_machine_ran_error(
    mock_set_arn: mock.MagicMock,
    mock_get_last_execution: mock.MagicMock,
    mock_time: mock.MagicMock,
) -> None:
    mock_set_arn.return_value = 'testArn'
    mock_time.return_value = None
    arn = 'newArn'
    mock_get_last_execution.return_value = arn
    handler = step_function_handler.StepFunctionHandler('test')
    handler.last_execution = arn
    with pytest.raises(AssertionError):
        handler.assert_machine_ran()


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.assert_execution_success'
)
@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_execute_with_all_params(
    mock_session_class: mock.MagicMock,
    mock_set_arn: mock.MagicMock,
    mock_assert_execution_success: mock.MagicMock,
) -> None:
    mock_set_arn.return_value = 'testArn'
    response = {'executionArn': 'testExecution'}
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.start_execution.return_value = response
    mock_assert_execution_success.return_value = None
    handler = step_function_handler.StepFunctionHandler('test')
    input_data = {'key': 'value'}
    handler.execute(
        input_data,
        True,
        timeout=90,
        execution_name='customize-sf-execution-name',
    )
    mock_client.start_execution.assert_called_with(
        stateMachineArn='testArn',
        input='{"key": "value"}',
        name='customize-sf-execution-name',
    )
    mock_assert_execution_success.assert_called_with('testExecution', 90)


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_execute_passes_kwargs_through(
    mock_session_class: mock.MagicMock, mock_set_arn: mock.MagicMock
) -> None:
    mock_set_arn.return_value = 'testArn'
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.start_execution.return_value = {'executionArn': 'testExecution'}
    handler = step_function_handler.StepFunctionHandler('test')
    handler.execute({'key': 'value'}, False, traceHeader='trace-123')
    mock_client.start_execution.assert_called_with(
        stateMachineArn='testArn',
        input='{"key": "value"}',
        traceHeader='trace-123',
    )


@mock.patch(
    'aws_testing_utils.step_function_handler.StepFunctionHandler.set_machine_arn'
)
@mock.patch('boto3.session.Session')
def test_execute_reserved_keys_win_over_kwargs(
    mock_session_class: mock.MagicMock, mock_set_arn: mock.MagicMock
) -> None:
    # Callers can't override the resolved ARN/input, and execution_name wins
    # over a name passed in kwargs.
    mock_set_arn.return_value = 'testArn'
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.start_execution.return_value = {'executionArn': 'testExecution'}
    handler = step_function_handler.StepFunctionHandler('test')
    handler.execute(
        {'key': 'value'},
        False,
        execution_name='real',
        stateMachineArn='hacker',
        input='evil',
        name='override',
    )
    mock_client.start_execution.assert_called_with(
        stateMachineArn='testArn',
        input='{"key": "value"}',
        name='real',
    )