moto/tests/test_stepfunctions/parser/__init__.py

import json
import os
from functools import wraps
from time import sleep
from typing import TYPE_CHECKING, Callable, TypeVar
from uuid import uuid4

import boto3
import requests

from moto import mock_aws, settings
from tests import DEFAULT_ACCOUNT_ID as ACCOUNT_ID
from tests.test_stepfunctions.parser.templates.templates import load_template

if TYPE_CHECKING:
    from typing_extensions import ParamSpec

    P = ParamSpec("P")

T = TypeVar("T")


def aws_verified(func):
    """
    Function that is verified to work against AWS.
    Can be run against AWS at any time by setting:
      MOTO_TEST_ALLOW_AWS_REQUEST=true

    If this environment variable is not set, the function runs in a `mock_aws` context.
    """

    @wraps(func)
    def pagination_wrapper():
        allow_aws_request = (
            os.environ.get("MOTO_TEST_ALLOW_AWS_REQUEST", "false").lower() == "true"
        )

        if allow_aws_request:
            return func()
        else:
            with mock_aws():
                requests.post(
                    f"http://{base_url}/moto-api/config",
                    json={"stepfunctions": {"execute_state_machine": True}},
                )
                resp = func()
                requests.post(
                    f"http://{base_url}/moto-api/config",
                    json={"stepfunctions": {"execute_state_machine": False}},
                )
                return resp

    return pagination_wrapper


def verify_execution_result(
    _verify_result, expected_status, tmpl_name, exec_input=None, sleep_time=0
):
    iam = boto3.client("iam", region_name="us-east-1")
    role_name = f"sfn_role_{str(uuid4())[0:6]}"
    sfn_role = iam.create_role(
        RoleName=role_name,
        AssumeRolePolicyDocument=json.dumps(sfn_role_policy),
        Path="/",
    )["Role"]["Arn"]
    iam.put_role_policy(
        PolicyDocument=json.dumps(sfn_allow_dynamodb),
        PolicyName="allowLambdaInvoke",
        RoleName=role_name,
    )
    sleep(sleep_time)

    client = boto3.client("stepfunctions", region_name="us-east-1")
    execution_arn, state_machine_arn = _start_execution(
        client, load_template(tmpl_name), exec_input, sfn_role
    )
    for _ in range(10):
        execution = client.describe_execution(executionArn=execution_arn)
        if execution["status"] == expected_status:
            result = _verify_result(client, execution, execution_arn)
            if result is not False:
                client.delete_state_machine(stateMachineArn=state_machine_arn)
                iam.delete_role_policy(
                    RoleName=role_name, PolicyName="allowLambdaInvoke"
                )
                iam.delete_role(RoleName=role_name)
                break
        sleep(0.1)
    else:
        client.delete_state_machine(stateMachineArn=state_machine_arn)
        iam.delete_role_policy(RoleName=role_name, PolicyName="allowLambdaInvoke")
        iam.delete_role(RoleName=role_name)
        assert False, "Should have failed already"


def _start_execution(client, definition, exec_input, sfn_role):
    name = "sfn_name"
    #
    response = client.create_state_machine(
        name=name, definition=json.dumps(definition), roleArn=sfn_role
    )
    state_machine_arn = response["stateMachineArn"]
    execution = client.start_execution(
        name="exec1", stateMachineArn=state_machine_arn, input=exec_input or "{}"
    )
    execution_arn = execution["executionArn"]
    return execution_arn, state_machine_arn


def _get_default_role():
    return "arn:aws:iam::" + ACCOUNT_ID + ":role/unknown_sf_role"


base_url = "localhost:5000" if settings.TEST_SERVER_MODE else "motoapi.amazonaws.com"


def enable_sfn_parser(func: "Callable[P, T]") -> "Callable[P, T]":
    @wraps(func)
    def pagination_wrapper(*args: "P.args", **kwargs: "P.kwargs") -> T:  # type: ignore
        requests.post(
            f"http://{base_url}/moto-api/config",
            json={"stepfunctions": {"execute_state_machine": True}},
        )
        try:
            res = func(*args, **kwargs)
        finally:
            requests.post(
                f"http://{base_url}/moto-api/config",
                json={"stepfunctions": {"execute_state_machine": False}},
            )
        return res

    return pagination_wrapper


sfn_role_policy = {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {"Service": "states.amazonaws.com"},
            "Action": "sts:AssumeRole",
        }
    ],
}

sfn_allow_lambda_invoke = {
    "Version": "2012-10-17",
    "Statement": [
        {"Effect": "Allow", "Action": ["lambda:InvokeFunction"], "Resource": ["*"]}
    ],
}

sfn_allow_dynamodb = {
    "Version": "2012-10-17",
    "Statement": [{"Effect": "Allow", "Action": ["dynamodb:*"], "Resource": ["*"]}],
}