localstack-core 4.10.1.dev7__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
 
@@ -103,6 +104,14 @@ def create_subscription_arn(topic_arn: str) -> str:
     return f"{topic_arn}:{uuid4()}"
 
 
+def create_platform_endpoint_arn(
+    platform_application_arn: str,
+) -> str:
+    # This is the format of an Endpoint Arn
+    # arn:aws:sns:us-west-2:1234567890:endpoint/GCM/MyApplication/12345678-abcd-9012-efgh-345678901234
+    return f"{platform_application_arn.replace('app', 'endpoint', 1)}/{uuid4()}"
+
+
 def encode_subscription_token_with_region(region: str) -> str:
     """
     Create a 64 characters Subscription Token with the region encoded
@@ -128,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/sqs/models.py

@@ -314,7 +314,8 @@ class SqsQueue:
     purge_timestamp: float | None
 
     delayed: set[SqsMessage]
-    inflight: set[SqsMessage]
+    # Simulating an ordered set in python. Only the keys are used and of interest.
+    inflight: dict[SqsMessage, None]
     receipts: dict[str, SqsMessage]
 
     def __init__(self, name: str, region: str, account_id: str, attributes=None, tags=None) -> None:
@@ -326,7 +327,7 @@ class SqsQueue:
         self.tags = tags or {}
 
         self.delayed = set()
-        self.inflight = set()
+        self.inflight = {}
         self.receipts = {}
 
         self.attributes = self.default_attributes()
@@ -513,7 +514,7 @@ class SqsQueue:
                 )
                 # Terminating the visibility timeout for a message
                 # https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-visibility-timeout.html#terminating-message-visibility-timeout
-                self.inflight.remove(standard_message)
+                del self.inflight[standard_message]
                 self._put_message(standard_message)
 
     def remove(self, receipt_handle: str):
@@ -606,9 +607,17 @@ class SqsQueue:
                     standard_message,
                     self.arn,
                 )
-                self.inflight.remove(standard_message)
+                del self.inflight[standard_message]
                 self._put_message(standard_message)
 
+    def add_inflight_message(self, message: SqsMessage):
+        """
+        We are simulating an ordered set with a dict. When a value is added, it is added as key to the dict, which
+        is all we need. Hence all "values" in this ordered set are None
+        :param message: The message to put in flight
+        """
+        self.inflight[message] = None
+
     def enqueue_delayed_messages(self):
         if not self.delayed:
             return
@@ -779,7 +788,6 @@ class SqsQueue:
 
 class StandardQueue(SqsQueue):
     visible: InterruptiblePriorityQueue[SqsMessage]
-    inflight: set[SqsMessage]
 
     def __init__(self, name: str, region: str, account_id: str, attributes=None, tags=None) -> None:
         super().__init__(name, region, account_id, attributes, tags)
@@ -923,13 +931,13 @@ class StandardQueue(SqsQueue):
             if message.visibility_timeout == 0:
                 self.visible.put_nowait(message)
             else:
-                self.inflight.add(message)
+                self.add_inflight_message(message)
 
         return result
 
     def _on_remove_message(self, message: SqsMessage):
         try:
-            self.inflight.remove(message)
+            del self.inflight[message]
         except KeyError:
             # this likely means the message was removed with an expired receipt handle unfortunately this
             # means we need to scan the queue for the element and remove it from there, and then re-heapify
@@ -1149,6 +1157,26 @@ class FifoQueue(SqsQueue):
             elif previously_empty:
                 self.message_group_queue.put_nowait(message_group)
 
+    def requeue_inflight_messages(self):
+        if not self.inflight:
+            return
+
+        with self.mutex:
+            messages = list(self.inflight)
+            for standard_message in messages:
+                # in fifo, an invisible message blocks potentially visible messages afterwards
+                # this can happen for example if multiple message of the same group are received at once, then one
+                # message of this batch has its visibility timeout extended
+                if not standard_message.is_visible:
+                    return
+                LOG.debug(
+                    "re-queueing inflight messages %s into queue %s",
+                    standard_message,
+                    self.arn,
+                )
+                del self.inflight[standard_message]
+                self._put_message(standard_message)
+
     def remove_expired_messages(self):
         with self.mutex:
             retention_period = self.message_retention_period
@@ -1278,8 +1306,7 @@ class FifoQueue(SqsQueue):
                 if message.visibility_timeout == 0:
                     self._put_message(message)
                 else:
-                    self.inflight.add(message)
-
+                    self.add_inflight_message(message)
         return result
 
     def _on_remove_message(self, message: SqsMessage):
@@ -1288,7 +1315,7 @@ class FifoQueue(SqsQueue):
 
         with self.mutex:
             try:
-                self.inflight.remove(message)
+                del self.inflight[message]
             except KeyError:
                 # in FIFO queues, this should not happen, as expired receipt handles cannot be used to
                 # delete a message.

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