import time

from botocore.stub import Stubber
import pytest

import main


@pytest.fixture
def mock_state_machine_arn():
    return "arn:aws:states:us-east-1:123456789012:execution:my-state-machine"


@pytest.fixture
def mock_execution_arn(mock_state_machine_arn):
    return f"${mock_state_machine_arn}:my-execution"


@pytest.fixture
def step_functions_stubber():
    """Return stubber for the Step Functions client."""
    with Stubber(main.stepfunctions_client) as stubber:
        yield stubber


def test_start_state_machine(
    step_functions_stubber, mock_state_machine_arn, mock_execution_arn
):
    step_functions_stubber.add_response(
        "start_execution",
        {"executionArn": mock_execution_arn, "startDate": time.time()},
        {"stateMachineArn": mock_state_machine_arn},
    )

    execution_arn = main.start_state_machine(mock_state_machine_arn)
    assert execution_arn == mock_execution_arn


def test_start_state_machine_with_optional_args(
    step_functions_stubber, mock_state_machine_arn, mock_execution_arn
):
    step_functions_stubber.add_response(
        "start_execution",
        {"executionArn": mock_execution_arn, "startDate": time.time()},
        {
            "stateMachineArn": mock_state_machine_arn,
            "name": "my-execution-name",
            "input": '{"foo":"bar"}',
        },
    )

    execution_arn = main.start_state_machine(
        mock_state_machine_arn, "my-execution-name", '{"foo":"bar"}'
    )
    assert execution_arn == mock_execution_arn


def test_wait_for_state_machine_to_complete_success(
    step_functions_stubber, mock_state_machine_arn, mock_execution_arn
):
    step_functions_stubber.add_response(
        "describe_execution",
        {
            "executionArn": mock_execution_arn,
            "stateMachineArn": mock_state_machine_arn,
            "startDate": time.time(),
            "status": "RUNNING",
        },
    )
    step_functions_stubber.add_response(
        "describe_execution",
        {
            "executionArn": mock_execution_arn,
            "stateMachineArn": mock_state_machine_arn,
            "startDate": time.time(),
            "status": "SUCCEEDED",
        },
    )

    main.wait_for_state_machine_to_complete(mock_execution_arn, 0)

    step_functions_stubber.assert_no_pending_responses()


def test_wait_for_state_machine_to_complete_aborted(
    step_functions_stubber, mock_state_machine_arn, mock_execution_arn
):
    step_functions_stubber.add_response(
        "describe_execution",
        {
            "executionArn": mock_execution_arn,
            "stateMachineArn": mock_state_machine_arn,
            "startDate": time.time(),
            "status": "ABORTED",
        },
    )

    with pytest.raises(Exception) as e:
        main.wait_for_state_machine_to_complete(mock_execution_arn, 0)
        assert str(e) == "State machine did not complete successfully. Status: ABORTED."


def test_wait_for_state_machine_to_complete_failure(
    step_functions_stubber, mock_state_machine_arn, mock_execution_arn
):
    step_functions_stubber.add_response(
        "describe_execution",
        {
            "cause": "just cause",
            "error": "failed",
            "executionArn": mock_execution_arn,
            "stateMachineArn": mock_state_machine_arn,
            "startDate": time.time(),
            "status": "FAILED",
        },
    )

    with pytest.raises(Exception) as e:
        main.wait_for_state_machine_to_complete(mock_execution_arn, 0)
        assert (
            str(e)
            == "State machine did not complete successfully. Status: ABORTED.\nError: failed\nCause: just cause"
        )