localstack-core 4.10.1.dev42__py3-none-any.whl → 4.11.2.dev14__py3-none-any.whl

localstack/services/sns/v2/utils.py

@@ -6,6 +6,7 @@ from botocore.utils import InvalidArnException
 
 from localstack.aws.api.sns import InvalidParameterException
 from localstack.services.sns.constants import E164_REGEX, VALID_SUBSCRIPTION_ATTR_NAME
+from localstack.services.sns.v2.models import SnsStore, SnsSubscription
 from localstack.utils.aws.arns import ArnData, parse_arn
 from localstack.utils.strings import short_uid, to_bytes, to_str
 
@@ -136,3 +137,14 @@ def get_region_from_subscription_token(token: str) -> str:
         return bytes.fromhex(region).decode("utf-8")
     except (IndexError, ValueError, TypeError, UnicodeDecodeError):
         raise InvalidParameterException("Invalid parameter: Token")
+
+
+def get_topic_subscriptions(store: SnsStore, topic_arn: str) -> list[SnsSubscription]:
+    # TODO: delete this once the legacy v1 implementation has been removed
+    if hasattr(store, "topic_subscriptions"):
+        sub_arns = store.topic_subscriptions.get(topic_arn, [])
+    else:
+        sub_arns: list[str] = store.topics[topic_arn].get("subscriptions", [])
+
+    subscriptions = [store.subscriptions[k] for k in sub_arns if k in store.subscriptions]
+    return subscriptions

localstack/services/stepfunctions/asl/component/common/path/result_path.py

@@ -18,7 +18,7 @@ class ResultPath(EvalComponent):
     def _eval_body(self, env: Environment) -> None:
         state_input = env.states.get_input()
 
-        # Discard task output if there is one, and set the output ot be the state's input.
+        # Discard task output if there is one, and set the output to be the state's input.
         if self.result_path_src is None:
             env.stack.clear()
             env.stack.append(state_input)

localstack/services/stepfunctions/asl/component/state/state_execution/execute_state.py

@@ -251,7 +251,6 @@ class ExecutionState(CommonStateField, abc.ABC):
                 )
                 error_output = self._construct_error_output_value(failure_event=failure_event)
                 env.states.set_error_output(error_output)
-                env.states.set_result(error_output)
 
                 if self.retry:
                     retry_outcome: RetryOutcome = self._handle_retry(

localstack/services/stepfunctions/asl/component/state/state_execution/state_map/state_map.py

@@ -311,7 +311,6 @@ class StateMap(ExecutionState):
                 failure_event: FailureEvent = self._from_error(env=env, ex=ex)
                 error_output = self._construct_error_output_value(failure_event=failure_event)
                 env.states.set_error_output(error_output)
-                env.states.set_result(error_output)
 
                 if self.retry:
                     retry_outcome: RetryOutcome = self._handle_retry(

localstack/services/stepfunctions/asl/component/state/state_execution/state_task/lambda_eval_utils.py

@@ -6,13 +6,13 @@ from localstack.aws.api.lambda_ import InvocationResponse
 from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.credentials import (
     StateCredentials,
 )
-from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.mock_eval_utils import (
-    eval_mocked_response,
+from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.local_mock_eval_utils import (
+    eval_local_mocked_response,
 )
 from localstack.services.stepfunctions.asl.eval.environment import Environment
 from localstack.services.stepfunctions.asl.utils.boto_client import boto_client_for
 from localstack.services.stepfunctions.asl.utils.encoding import to_json_str
-from localstack.services.stepfunctions.mocking.mock_config import MockedResponse
+from localstack.services.stepfunctions.local_mocking.mock_config import LocalMockedResponse
 from localstack.utils.collections import select_from_typed_dict
 from localstack.utils.strings import to_bytes
 
@@ -42,9 +42,9 @@ def _from_payload(payload_streaming_body: IO[bytes]) -> Any | str:
         return decoded_data
 
 
-def _mocked_invoke_lambda_function(env: Environment) -> InvocationResponse:
-    mocked_response: MockedResponse = env.get_current_mocked_response()
-    eval_mocked_response(env=env, mocked_response=mocked_response)
+def _local_mocked_invoke_lambda_function(env: Environment) -> InvocationResponse:
+    mocked_response: LocalMockedResponse = env.get_current_local_mocked_response()
+    eval_local_mocked_response(env=env, mocked_response=mocked_response)
     invocation_resp: InvocationResponse = env.stack.pop()
     return invocation_resp
 
@@ -68,8 +68,8 @@ def _invoke_lambda_function(
 def execute_lambda_function_integration(
     env: Environment, parameters: dict, region: str, state_credentials: StateCredentials
 ) -> None:
-    if env.is_mocked_mode():
-        invocation_response: InvocationResponse = _mocked_invoke_lambda_function(env=env)
+    if env.is_local_mocked_mode():
+        invocation_response: InvocationResponse = _local_mocked_invoke_lambda_function(env=env)
     else:
         invocation_response: InvocationResponse = _invoke_lambda_function(
             parameters=parameters, region=region, state_credentials=state_credentials

localstack/services/stepfunctions/asl/component/state/state_execution/state_task/{mock_eval_utils.py → local_mock_eval_utils.py}

@@ -10,14 +10,16 @@ from localstack.services.stepfunctions.asl.component.common.error_name.failure_e
 )
 from localstack.services.stepfunctions.asl.eval.environment import Environment
 from localstack.services.stepfunctions.asl.eval.event.event_detail import EventDetails
-from localstack.services.stepfunctions.mocking.mock_config import (
-    MockedResponse,
-    MockedResponseReturn,
-    MockedResponseThrow,
+from localstack.services.stepfunctions.local_mocking.mock_config import (
+    LocalMockedResponse,
+    LocalMockedResponseReturn,
+    LocalMockedResponseThrow,
 )
 
 
-def _eval_mocked_response_throw(env: Environment, mocked_response: MockedResponseThrow) -> None:
+def _eval_mocked_response_throw(
+    env: Environment, mocked_response: LocalMockedResponseThrow
+) -> None:
     task_failed_event_details = TaskFailedEventDetails(
         error=mocked_response.error, cause=mocked_response.cause
     )
@@ -31,15 +33,17 @@ def _eval_mocked_response_throw(env: Environment, mocked_response: MockedRespons
     raise FailureEventException(failure_event=failure_event)
 
 
-def _eval_mocked_response_return(env: Environment, mocked_response: MockedResponseReturn) -> None:
+def _eval_mocked_response_return(
+    env: Environment, mocked_response: LocalMockedResponseReturn
+) -> None:
     payload_copy = copy.deepcopy(mocked_response.payload)
     env.stack.append(payload_copy)
 
 
-def eval_mocked_response(env: Environment, mocked_response: MockedResponse) -> None:
-    if isinstance(mocked_response, MockedResponseReturn):
+def eval_local_mocked_response(env: Environment, mocked_response: LocalMockedResponse) -> None:
+    if isinstance(mocked_response, LocalMockedResponseReturn):
         _eval_mocked_response_return(env=env, mocked_response=mocked_response)
-    elif isinstance(mocked_response, MockedResponseThrow):
+    elif isinstance(mocked_response, LocalMockedResponseThrow):
         _eval_mocked_response_throw(env=env, mocked_response=mocked_response)
     else:
         raise RuntimeError(f"Invalid MockedResponse type '{type(mocked_response)}'")

localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service.py

@@ -33,8 +33,8 @@ from localstack.services.stepfunctions.asl.component.common.error_name.states_er
 from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.credentials import (
     StateCredentials,
 )
-from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.mock_eval_utils import (
-    eval_mocked_response,
+from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.local_mock_eval_utils import (
+    eval_local_mocked_response,
 )
 from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.service.resource import (
     ResourceRuntimePart,
@@ -47,7 +47,7 @@ from localstack.services.stepfunctions.asl.component.state.state_props import St
 from localstack.services.stepfunctions.asl.eval.environment import Environment
 from localstack.services.stepfunctions.asl.eval.event.event_detail import EventDetails
 from localstack.services.stepfunctions.asl.utils.encoding import to_json_str
-from localstack.services.stepfunctions.mocking.mock_config import MockedResponse
+from localstack.services.stepfunctions.local_mocking.mock_config import LocalMockedResponse
 from localstack.services.stepfunctions.quotas import is_within_size_quota
 from localstack.utils.strings import camel_to_snake_case, snake_to_camel_case, to_bytes, to_str
 
@@ -356,9 +356,9 @@ class StateTaskService(StateTask, abc.ABC):
         normalised_parameters = copy.deepcopy(raw_parameters)
         self._normalise_parameters(normalised_parameters)
 
-        if env.is_mocked_mode():
-            mocked_response: MockedResponse = env.get_current_mocked_response()
-            eval_mocked_response(env=env, mocked_response=mocked_response)
+        if env.is_local_mocked_mode():
+            mocked_response: LocalMockedResponse = env.get_current_local_mocked_response()
+            eval_local_mocked_response(env=env, mocked_response=mocked_response)
         else:
             self._eval_service_task(
                 env=env,

localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_callback.py

@@ -346,7 +346,7 @@ class StateTaskServiceCallback(StateTaskService, abc.ABC):
                     )
                 ),
             )
-            if not env.is_mocked_mode():
+            if not env.is_local_mocked_mode() and not env.is_test_state_mocked_mode():
                 self._eval_integration_pattern(
                     env=env,
                     resource_runtime_part=resource_runtime_part,

localstack/services/stepfunctions/asl/component/state/state_fail/state_fail.py

@@ -7,6 +7,9 @@ from localstack.services.stepfunctions.asl.component.common.error_name.failure_e
     FailureEventException,
 )
 from localstack.services.stepfunctions.asl.component.state.state import CommonStateField
+from localstack.services.stepfunctions.asl.component.state.state_continue_with import (
+    ContinueWithEnd,
+)
 from localstack.services.stepfunctions.asl.component.state.state_fail.cause_decl import CauseDecl
 from localstack.services.stepfunctions.asl.component.state.state_fail.error_decl import ErrorDecl
 from localstack.services.stepfunctions.asl.component.state.state_props import StateProps
@@ -27,6 +30,7 @@ class StateFail(CommonStateField):
         super().from_state_props(state_props)
         self.cause = state_props.get(CauseDecl)
         self.error = state_props.get(ErrorDecl)
+        self.continue_with = ContinueWithEnd()
 
     def _eval_state(self, env: Environment) -> None:
         task_failed_event_details = TaskFailedEventDetails()

localstack/services/stepfunctions/asl/component/test_state/state/base_mock.py

@@ -0,0 +1,118 @@
+import abc
+import copy
+from typing import Generic, TypeVar
+
+from localstack.services.stepfunctions.asl.component.eval_component import EvalComponent
+from localstack.services.stepfunctions.asl.component.state.state import CommonStateField
+from localstack.services.stepfunctions.asl.component.state.state_continue_with import (
+    ContinueWithNext,
+)
+from localstack.services.stepfunctions.asl.eval.test_state.environment import TestStateEnvironment
+from localstack.services.stepfunctions.asl.utils.encoding import to_json_str
+from localstack.services.stepfunctions.backend.test_state.test_state_mock import (
+    TestStateResponseReturn,
+    TestStateResponseThrow,
+    eval_mocked_response_throw,
+)
+
+T = TypeVar("T", bound=CommonStateField)
+
+
+class MockedBaseState(Generic[T], abc.ABC):
+    is_single_state: bool
+    _wrapped: T
+
+    def __init__(self, wrapped: T):
+        super().__init__()
+        self._wrapped = wrapped
+        self.apply_patches()
+
+    def apply_patches(self):
+        self._apply_patches()
+
+        original_eval_body = self._wrapped._eval_body
+        self._wrapped._eval_body = self.wrap_with_post_return(
+            original_eval_body, self.stop_execution
+        )
+
+    @abc.abstractmethod
+    def _apply_patches(self): ...
+
+    @classmethod
+    def wrap(cls, state: T, is_single_state: bool = False) -> T:
+        cls.is_single_state = is_single_state
+        cls._wrapped = state
+        return cls(state)._wrapped
+
+    def __getattr__(self, attr: str):
+        return getattr(self._wrapped, attr)
+
+    @classmethod
+    def before_mock(self, env: TestStateEnvironment):
+        return
+
+    @classmethod
+    def do_mock(self, env: TestStateEnvironment):
+        mocked_response = env.mock.get_next_result()
+        if not mocked_response:
+            return
+
+        if isinstance(mocked_response, TestStateResponseThrow):
+            eval_mocked_response_throw(env, mocked_response)
+            return
+
+        if isinstance(mocked_response, TestStateResponseReturn):
+            result_copy = copy.deepcopy(mocked_response.payload)
+            env.stack.append(result_copy)
+
+    @classmethod
+    def after_mock(self, env: TestStateEnvironment):
+        return
+
+    @classmethod
+    def wrap_with_mock(cls, original_method):
+        def wrapper(env: TestStateEnvironment, *args, **kwargs):
+            if not env.mock.is_mocked():
+                original_method(env, *args, **kwargs)
+                return
+
+            cls.before_mock(env)
+            try:
+                cls.do_mock(env)
+            finally:
+                cls.after_mock(env)
+
+        return wrapper
+
+    @staticmethod
+    def wrap_with_post_return(method, post_return_fn):
+        def wrapper(env: TestStateEnvironment, *args, **kwargs):
+            try:
+                method(env, *args, **kwargs)
+            finally:
+                post_return_fn(env)
+
+        return wrapper
+
+    @staticmethod
+    def _eval_with_inspect(component: EvalComponent, key: str):
+        if not component:
+            return
+
+        eval_body_fn = component._eval_body
+
+        def _update(env: TestStateEnvironment, *args, **kwargs):
+            # if inspectionData already populated, don't execute again
+            if key in env.inspection_data:
+                return
+
+            eval_body_fn(env, *args, **kwargs)
+            result = env.stack[-1]
+            env.inspection_data[key] = to_json_str(result)
+
+        component._eval_body = MockedBaseState.wrap_with_post_return(eval_body_fn, _update)
+
+    def stop_execution(self, env: TestStateEnvironment):
+        if isinstance(self._wrapped.continue_with, ContinueWithNext):
+            if next_state := self._wrapped.continue_with.next_state:
+                env.set_choice_selected(next_state.name)

localstack/services/stepfunctions/asl/component/test_state/state/common.py

@@ -0,0 +1,82 @@
+from localstack.services.stepfunctions.asl.component.state.state import CommonStateField
+from localstack.services.stepfunctions.asl.component.state.state_choice.state_choice import (
+    StateChoice,
+)
+from localstack.services.stepfunctions.asl.component.state.state_continue_with import (
+    ContinueWithEnd,
+)
+from localstack.services.stepfunctions.asl.component.state.state_fail.state_fail import StateFail
+from localstack.services.stepfunctions.asl.component.state.state_pass.state_pass import StatePass
+from localstack.services.stepfunctions.asl.component.state.state_succeed.state_succeed import (
+    StateSucceed,
+)
+from localstack.services.stepfunctions.asl.component.test_state.state.base_mock import (
+    MockedBaseState,
+)
+from localstack.services.stepfunctions.asl.eval.test_state.environment import TestStateEnvironment
+
+
+class MockedCommonState(MockedBaseState[CommonStateField]):
+    def add_inspection_data(self, env: TestStateEnvironment):
+        state = self._wrapped
+
+        if not isinstance(state, StatePass):
+            if not self.is_single_state:
+                return
+
+            if "afterInputPath" not in env.inspection_data:
+                env.inspection_data["afterInputPath"] = env.states.get_input()
+            return
+
+        # If not a terminal state, only populate inspection data from pre-processor.
+        if not isinstance(self._wrapped.continue_with, ContinueWithEnd):
+            return
+
+        if state.result:
+            # TODO: investigate interactions between these inspectionData field types.
+            # i.e parity tests shows that if "Result" is defined, 'afterInputPath' and 'afterParameters'
+            # cannot be present in the inspection data.
+            env.inspection_data.pop("afterInputPath", None)
+            env.inspection_data.pop("afterParameters", None)
+
+            if "afterResultSelector" not in env.inspection_data:
+                env.inspection_data["afterResultSelector"] = state.result.result_obj
+
+            if "afterResultPath" not in env.inspection_data:
+                env.inspection_data["afterResultPath"] = env.inspection_data.get(
+                    "afterResultSelector", env.states.get_input()
+                )
+            return
+
+        if "afterInputPath" not in env.inspection_data:
+            env.inspection_data["afterInputPath"] = env.states.get_input()
+
+        if "afterParameters" not in env.inspection_data:
+            env.inspection_data["afterParameters"] = env.inspection_data.get(
+                "afterInputPath", env.states.get_input()
+            )
+
+        if "afterResultSelector" not in env.inspection_data:
+            env.inspection_data["afterResultSelector"] = env.inspection_data["afterParameters"]
+
+        if "afterResultPath" not in env.inspection_data:
+            env.inspection_data["afterResultPath"] = env.inspection_data.get(
+                "afterResultSelector", env.states.get_input()
+            )
+
+    def _apply_patches(self):
+        if not isinstance(self._wrapped, (StatePass, StateFail, StateChoice, StateSucceed)):
+            raise ValueError("Needs to be a Pass, Fail, Choice, or Succeed state.")
+
+        original_eval_body = self.wrap_with_mock(self._wrapped._eval_body)
+
+        def mock_eval_execution(env: TestStateEnvironment):
+            original_eval_body(env)
+            env.set_choice_selected(env.next_state_name)
+
+        mock_eval_execution = self.wrap_with_post_return(
+            method=mock_eval_execution,
+            post_return_fn=self.add_inspection_data,
+        )
+
+        self._wrapped._eval_body = mock_eval_execution

localstack/services/stepfunctions/asl/component/test_state/state/execution.py

@@ -0,0 +1,139 @@
+from collections.abc import Callable
+from functools import partial
+
+from localstack.services.stepfunctions.asl.component.common.catch.catcher_outcome import (
+    CatcherOutcomeCaught,
+)
+from localstack.services.stepfunctions.asl.component.common.error_name.failure_event import (
+    FailureEvent,
+)
+from localstack.services.stepfunctions.asl.component.common.query_language import (
+    QueryLanguageMode,
+)
+from localstack.services.stepfunctions.asl.component.common.retry.retrier_decl import RetrierDecl
+from localstack.services.stepfunctions.asl.component.common.retry.retrier_outcome import (
+    RetrierOutcome,
+)
+from localstack.services.stepfunctions.asl.component.common.retry.retry_outcome import RetryOutcome
+from localstack.services.stepfunctions.asl.component.state.state_execution.execute_state import (
+    ExecutionState,
+)
+from localstack.services.stepfunctions.asl.component.test_state.state.base_mock import (
+    MockedBaseState,
+)
+from localstack.services.stepfunctions.asl.eval.test_state.environment import TestStateEnvironment
+from localstack.services.stepfunctions.asl.utils.encoding import to_json_str
+
+
+class MockedStateExecution(MockedBaseState[ExecutionState]):
+    def add_inspection_data(self, env: TestStateEnvironment):
+        if self._wrapped.query_language.query_language_mode == QueryLanguageMode.JSONPath:
+            if "afterResultSelector" not in env.inspection_data:
+                # HACK: A DistributedItemProcessorEvalInput is added to the stack and never popped off
+                # during an error case. So we need to check the inspected value is correct before
+                # adding it to our inspectionData.
+                if isinstance(env.stack[-1], (dict, str, int, float)):
+                    env.inspection_data["afterResultSelector"] = to_json_str(env.stack[-1])
+
+        if catch := self._wrapped.catch:
+            for ind, catcher in enumerate(catch.catchers):
+                original_fn = catcher._eval_body
+                catcher._eval_body = self.with_catch_state_id(original_fn, ind)
+
+        if retry := self._wrapped.retry:
+            for ind, retrier in enumerate(retry.retriers):
+                original_fn = retrier._eval_body
+                retrier._eval_body = self.with_retry_state_id(retrier, ind)
+
+    def _apply_patches(self):
+        if not isinstance(self._wrapped, ExecutionState):
+            raise ValueError("Can only apply MockedStateExecution patches to an ExecutionState")
+        state = self._wrapped
+
+        if state.query_language.query_language_mode == QueryLanguageMode.JSONPath:
+            self._eval_with_inspect(self._wrapped.input_path, "afterInputPath")
+            self._eval_with_inspect(self._wrapped.result_path, "afterResultPath")
+
+        self._eval_with_inspect(self._wrapped.result_selector, "afterResultSelector")
+        original_eval_execution = self._wrapped._eval_execution
+
+        if self._wrapped.catch:
+            original_fn = self._wrapped._handle_catch
+            self._wrapped._handle_catch = partial(self._handle_catch, original_fn)
+
+        if self._wrapped.retry:
+            original_fn = self._wrapped._handle_retry
+            self._wrapped._handle_retry = partial(self._handle_retry, original_fn)
+
+        self._wrapped._eval_execution = self.wrap_with_post_return(
+            method=original_eval_execution,
+            post_return_fn=self.add_inspection_data,
+        )
+
+    @staticmethod
+    def with_catch_state_id(
+        original_eval_body: Callable[[TestStateEnvironment], None], state_id: int
+    ) -> Callable[[TestStateEnvironment], None]:
+        def _wrapped(env: TestStateEnvironment):
+            original_eval_body(env)
+
+            if isinstance(env.stack[-1], CatcherOutcomeCaught):
+                if not (error_details := env.inspection_data.get("errorDetails")):
+                    error_details = env.inspection_data["errorDetails"] = {}
+
+                error_details["catchIndex"] = state_id
+
+        return _wrapped
+
+    @staticmethod
+    def with_retry_state_id(
+        retrier: RetrierDecl, state_id: int
+    ) -> Callable[[TestStateEnvironment], None]:
+        original_retrier_eval_body = retrier._eval_body
+
+        def _wrapped(env: TestStateEnvironment):
+            if (retry_count := env.mock._state_configuration.get("retrierRetryCount", 0)) > 0:
+                retrier.max_attempts._store_attempt_number(env, retry_count - 1)
+
+            original_retrier_eval_body(env)
+
+            if not (error_details := env.inspection_data.get("errorDetails")):
+                error_details = env.inspection_data["errorDetails"] = {}
+
+            error_details["retryIndex"] = state_id
+            if env.stack[-1] == RetrierOutcome.Executed:
+                # TODO(gregfurman): Ideally, retryBackoffIntervalSeconds should be written to inspectionData
+                # within the retrier.backoff_rate decleration (perhaps at _access_next_multiplier).
+                rate = retrier.backoff_rate.rate
+                interval = retrier.interval_seconds.seconds
+                error_details["retryBackoffIntervalSeconds"] = int(interval * (rate**retry_count))
+
+        return _wrapped
+
+    @staticmethod
+    def _handle_catch(
+        original_handle_catch: Callable[[TestStateEnvironment, FailureEvent], None],
+        env: TestStateEnvironment,
+        failure_event: FailureEvent,
+    ) -> None:
+        original_handle_catch(env, failure_event)
+
+        spec: dict[str, str] = ExecutionState._construct_error_output_value(failure_event)
+        error, cause = spec.get("Error"), spec.get("Cause")
+
+        env.set_caught_error(env.next_state_name, error, cause)
+
+    @staticmethod
+    def _handle_retry(
+        original_handle_retry: Callable[[TestStateEnvironment, FailureEvent], RetryOutcome],
+        env: TestStateEnvironment,
+        failure_event: FailureEvent,
+    ) -> RetryOutcome:
+        res = original_handle_retry(env, failure_event)
+
+        spec: dict[str, str] = ExecutionState._construct_error_output_value(failure_event)
+        error, cause = spec.get("Error"), spec.get("Cause")
+
+        if res == RetryOutcome.CanRetry:
+            env.set_retriable_error(error, cause)
+        return res

localstack/services/stepfunctions/asl/component/test_state/state/map.py

@@ -0,0 +1,77 @@
+from localstack.services.stepfunctions.asl.component.common.error_name.states_error_name_type import (
+    StatesErrorNameType,
+)
+from localstack.services.stepfunctions.asl.component.common.query_language import (
+    QueryLanguageMode,
+)
+from localstack.services.stepfunctions.asl.component.state.state_execution.state_map.state_map import (
+    StateMap,
+)
+from localstack.services.stepfunctions.asl.component.test_state.state.base_mock import (
+    MockedBaseState,
+)
+from localstack.services.stepfunctions.asl.component.test_state.state.execution import (
+    MockedStateExecution,
+)
+from localstack.services.stepfunctions.asl.eval.test_state.environment import TestStateEnvironment
+from localstack.services.stepfunctions.backend.test_state.test_state_mock import (
+    TestStateResponseThrow,
+)
+
+
+class MockedStateMap(MockedBaseState[StateMap]):
+    def add_inspection_data(self, env: TestStateEnvironment):
+        if tolerated_failure_percentage := env.inspection_data.get("toleratedFailurePercentage"):
+            env.inspection_data["toleratedFailurePercentage"] = float(tolerated_failure_percentage)
+
+        if tolerated_failure_count := env.inspection_data.get("toleratedFailureCount"):
+            env.inspection_data["toleratedFailureCount"] = int(tolerated_failure_count)
+
+    @classmethod
+    def before_mock(cls, env: TestStateEnvironment):
+        if not env.mock or not env.mock._state_configuration:
+            return
+
+        if not cls._wrapped.catch and not cls._wrapped.retry:
+            return
+
+        if failure_count := env.mock._state_configuration.get("mapIterationFailureCount"):
+            max_failure_count = (
+                cls._wrapped.tolerated_failure_count_decl._eval_tolerated_failure_count(env)
+            )
+            if failure_count > max_failure_count:
+                error_response = TestStateResponseThrow(
+                    error=StatesErrorNameType.StatesExceedToleratedFailureThreshold.to_name(),
+                    cause="The specified tolerated failure threshold was exceeded",
+                )
+                env.mock.add_result(error_response)
+                return
+
+    def _apply_patches(self):
+        self._wrapped = MockedStateExecution.wrap(self._wrapped)
+
+        if self._wrapped.query_language.query_language_mode == QueryLanguageMode.JSONPath:
+            self._eval_with_inspect(self._wrapped.items_path, "afterInputPath")
+            self._eval_with_inspect(self._wrapped.item_selector, "afterItemsSelector")
+
+        original_eval_max_concurrency = self._wrapped.max_concurrency_decl._eval_max_concurrency
+        original_iteration_component_eval_body = self._wrapped.iteration_component._eval_body
+        original_eval_execution = self._wrapped._eval_execution
+
+        # HACK(gregfurman): Ideally we should be using the "$$.Map.Item.Index" to access each item of the
+        # mocked result list. This is turning out to be quite complicated, so instead just patch the
+        # StateMap's max concurrency decleration to always eval to '1' -- making the map run in serial.
+        def mock_max_concurrency(env: TestStateEnvironment) -> int:
+            # always set concurrency to 1 but inspection data is accurate to original
+            env.inspection_data["maxConcurrency"] = original_eval_max_concurrency(env)
+            return 1
+
+        self._wrapped._eval_execution = self.wrap_with_post_return(
+            method=original_eval_execution,
+            post_return_fn=self.add_inspection_data,
+        )
+
+        self._wrapped.max_concurrency_decl._eval_max_concurrency = mock_max_concurrency
+        self._wrapped.iteration_component._eval_body = self.wrap_with_mock(
+            original_iteration_component_eval_body
+        )

localstack/services/stepfunctions/asl/component/test_state/state/task.py

@@ -0,0 +1,44 @@
+from localstack.services.stepfunctions.asl.component.common.query_language import (
+    QueryLanguageMode,
+)
+from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.service.state_task_service import (
+    StateTaskService,
+)
+from localstack.services.stepfunctions.asl.component.state.state_execution.state_task.state_task import (
+    StateTask,
+)
+from localstack.services.stepfunctions.asl.component.test_state.state.base_mock import (
+    MockedBaseState,
+)
+from localstack.services.stepfunctions.asl.component.test_state.state.execution import (
+    MockedStateExecution,
+)
+from localstack.services.stepfunctions.asl.eval.test_state.environment import TestStateEnvironment
+from localstack.services.stepfunctions.asl.utils.encoding import to_json_str
+
+
+class MockedStateTask(MockedBaseState[StateTask]):
+    def add_inspection_data(self, env: TestStateEnvironment):
+        if self._wrapped.query_language.query_language_mode == QueryLanguageMode.JSONPath:
+            if "afterParameters" not in env.inspection_data:
+                env.inspection_data["afterParameters"] = to_json_str(env.states.get_input())
+
+    def _apply_patches(self):
+        self._wrapped = MockedStateExecution.wrap(self._wrapped)
+
+        if self._wrapped.query_language.query_language_mode == QueryLanguageMode.JSONPath:
+            self._eval_with_inspect(self._wrapped.parargs, "afterParameters")
+
+        if isinstance(self._wrapped, StateTaskService):
+            self._wrapped._eval_service_task = self.wrap_with_mock(self._wrapped._eval_service_task)
+
+        original_eval_execution = self._wrapped._eval_execution
+
+        def mock_eval_execution(env: TestStateEnvironment, *args, **kwargs):
+            original_eval_execution(env, *args, **kwargs)
+            result = to_json_str(env.stack[-1])
+            env.inspection_data["result"] = result
+
+        self._wrapped._eval_execution = self.wrap_with_post_return(
+            mock_eval_execution, self.add_inspection_data
+        )