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

localstack/services/stepfunctions/backend/test_state/execution.py

@@ -19,7 +19,9 @@ from localstack.services.stepfunctions.asl.eval.program_state import (
     ProgramState,
 )
 from localstack.services.stepfunctions.asl.eval.test_state.program_state import (
+    ProgramCaughtError,
     ProgramChoiceSelected,
+    ProgramRetriable,
 )
 from localstack.services.stepfunctions.asl.utils.encoding import to_json_str
 from localstack.services.stepfunctions.backend.activity import Activity
@@ -31,6 +33,7 @@ from localstack.services.stepfunctions.backend.state_machine import StateMachine
 from localstack.services.stepfunctions.backend.test_state.execution_worker import (
     TestStateExecutionWorker,
 )
+from localstack.services.stepfunctions.backend.test_state.test_state_mock import TestStateMock
 
 LOG = logging.getLogger(__name__)
 
@@ -38,6 +41,8 @@ LOG = logging.getLogger(__name__)
 class TestStateExecution(Execution):
     exec_worker: TestStateExecutionWorker | None
     next_state: str | None
+    state_name: str | None
+    mock: TestStateMock | None
 
     class TestCaseExecutionWorkerCommunication(BaseExecutionWorkerCommunication):
         _execution: TestStateExecution
@@ -48,6 +53,16 @@ class TestStateExecution(Execution):
                 self.execution.exec_status = ExecutionStatus.SUCCEEDED
                 self.execution.output = self.execution.exec_worker.env.states.get_input()
                 self.execution.next_state = exit_program_state.next_state_name
+            elif isinstance(exit_program_state, ProgramCaughtError):
+                self.execution.exec_status = ExecutionStatus.SUCCEEDED
+                self.execution.error = exit_program_state.error
+                self.execution.cause = exit_program_state.cause
+                self.execution.output = self.execution.exec_worker.env.states.get_input()
+                self.execution.next_state = exit_program_state.next_state_name
+            elif isinstance(exit_program_state, ProgramRetriable):
+                self.execution.exec_status = ExecutionStatus.SUCCEEDED
+                self.execution.error = exit_program_state.error
+                self.execution.cause = exit_program_state.cause
             else:
                 self._reflect_execution_status()
 
@@ -61,7 +76,9 @@ class TestStateExecution(Execution):
         state_machine: StateMachineInstance,
         start_date: Timestamp,
         activity_store: dict[Arn, Activity],
+        state_name: str | None = None,
         input_data: dict | None = None,
+        mock: TestStateMock | None = None,
     ):
         super().__init__(
             name=name,
@@ -79,6 +96,8 @@ class TestStateExecution(Execution):
         )
         self._execution_terminated_event = threading.Event()
         self.next_state = None
+        self.state_name = state_name
+        self.mock = mock
 
     def _get_start_execution_worker_comm(self) -> BaseExecutionWorkerCommunication:
         return self.TestCaseExecutionWorkerCommunication(self)
@@ -93,6 +112,8 @@ class TestStateExecution(Execution):
             exec_comm=self._get_start_execution_worker_comm(),
             cloud_watch_logging_session=self._cloud_watch_logging_session,
             activity_store=self._activity_store,
+            state_name=self.state_name,
+            mock=self.mock,
         )
 
     def publish_execution_status_change_event(self):
@@ -117,6 +138,21 @@ class TestStateExecution(Execution):
             test_state_output = TestStateOutput(
                 status=TestExecutionStatus.SUCCEEDED, nextState=self.next_state, output=output_str
             )
+        elif isinstance(exit_program_state, ProgramCaughtError):
+            output_str = to_json_str(self.output)
+            test_state_output = TestStateOutput(
+                status=TestExecutionStatus.CAUGHT_ERROR,
+                nextState=self.next_state,
+                output=output_str,
+                error=exit_program_state.error,
+                cause=exit_program_state.cause,
+            )
+        elif isinstance(exit_program_state, ProgramRetriable):
+            test_state_output = TestStateOutput(
+                status=TestExecutionStatus.RETRIABLE,
+                error=exit_program_state.error,
+                cause=exit_program_state.cause,
+            )
         else:
             # TODO: handle other statuses
             LOG.warning(

localstack/services/stepfunctions/backend/test_state/execution_worker.py

@@ -1,8 +1,13 @@
+from localstack.aws.api.stepfunctions import Arn, StateName
 from localstack.services.stepfunctions.asl.component.eval_component import EvalComponent
 from localstack.services.stepfunctions.asl.eval.environment import Environment
+from localstack.services.stepfunctions.asl.eval.evaluation_details import EvaluationDetails
 from localstack.services.stepfunctions.asl.eval.event.event_manager import (
     EventHistoryContext,
 )
+from localstack.services.stepfunctions.asl.eval.event.logging import (
+    CloudWatchLoggingSession,
+)
 from localstack.services.stepfunctions.asl.eval.states import (
     ContextObjectData,
     ExecutionData,
@@ -12,15 +17,41 @@ from localstack.services.stepfunctions.asl.eval.test_state.environment import Te
 from localstack.services.stepfunctions.asl.parse.test_state.asl_parser import (
     TestStateAmazonStateLanguageParser,
 )
+from localstack.services.stepfunctions.backend.activity import Activity
 from localstack.services.stepfunctions.backend.execution_worker import SyncExecutionWorker
+from localstack.services.stepfunctions.backend.execution_worker_comm import (
+    ExecutionWorkerCommunication,
+)
+from localstack.services.stepfunctions.backend.test_state.test_state_mock import TestStateMock
 
 
 class TestStateExecutionWorker(SyncExecutionWorker):
     env: TestStateEnvironment | None
+    state_name: str | None = None
+    mock: TestStateMock | None
+
+    def __init__(
+        self,
+        evaluation_details: EvaluationDetails,
+        exec_comm: ExecutionWorkerCommunication,
+        cloud_watch_logging_session: CloudWatchLoggingSession | None,
+        activity_store: dict[Arn, Activity],
+        state_name: StateName | None = None,
+        mock: TestStateMock | None = None,
+    ):
+        super().__init__(
+            evaluation_details,
+            exec_comm,
+            cloud_watch_logging_session,
+            activity_store,
+            local_mock_test_case=None,  # local mock is only applicable to SFN Local, but not for TestState
+        )
+        self.state_name = state_name
+        self.mock = mock
 
     def _get_evaluation_entrypoint(self) -> EvalComponent:
         return TestStateAmazonStateLanguageParser.parse(
-            self._evaluation_details.state_machine_details.definition
+            self._evaluation_details.state_machine_details.definition, self.state_name
         )[0]
 
     def _get_evaluation_environment(self) -> Environment:
@@ -43,4 +74,5 @@ class TestStateExecutionWorker(SyncExecutionWorker):
             event_history_context=EventHistoryContext.of_program_start(),
             cloud_watch_logging_session=self._cloud_watch_logging_session,
             activity_store=self._activity_store,
+            mock=self.mock,
         )

localstack/services/stepfunctions/backend/test_state/test_state_mock.py

@@ -0,0 +1,127 @@
+import copy
+import json
+from typing import Final
+
+from pydantic import (
+    ValidationError,
+)
+
+from localstack.aws.api.stepfunctions import (
+    HistoryEventType,
+    MockInput,
+    TaskFailedEventDetails,
+    TestStateConfiguration,
+)
+from localstack.services.stepfunctions.asl.component.common.error_name.error_name import ErrorName
+from localstack.services.stepfunctions.asl.component.common.error_name.failure_event import (
+    FailureEvent,
+    FailureEventException,
+)
+from localstack.services.stepfunctions.asl.component.state.state_type import StateType
+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.eval.states import (
+    ContextObjectData,
+)
+from localstack.services.stepfunctions.test_state.mock_config import (
+    TestStateContextObjectValidator,
+    TestStateMockedResponse,
+    TestStateResponseReturn,
+    TestStateResponseThrow,
+)
+
+
+def eval_mocked_response_throw(env: Environment, mocked_response: TestStateResponseThrow) -> None:
+    task_failed_event_details = TaskFailedEventDetails(
+        error=mocked_response.error, cause=mocked_response.cause
+    )
+    error_name = ErrorName(mocked_response.error)
+    failure_event = FailureEvent(
+        env=env,
+        error_name=error_name,
+        event_type=HistoryEventType.TaskFailed,  # TODO(gregfurman): Should this be state specific?
+        event_details=EventDetails(taskFailedEventDetails=task_failed_event_details),
+    )
+    raise FailureEventException(failure_event=failure_event)
+
+
+class TestStateMock:
+    _mock_input: MockInput | None
+    _state_configuration: TestStateConfiguration | None
+    _result_stack: Final[list[TestStateMockedResponse]]
+    _context: Final[ContextObjectData | None]
+
+    def __init__(
+        self,
+        mock_input: MockInput | None,
+        state_configuration: TestStateConfiguration | None,
+        context: str | None,
+    ):
+        self._mock_input = mock_input
+        self._state_configuration = state_configuration
+        self._result_stack = []
+        self._context = None
+
+        if not mock_input:
+            return
+
+        self._context = None if context is None else self.parse_context(context)
+
+        if mock_result_raw := mock_input.get("result"):
+            mock = json.loads(mock_result_raw)
+            self._result_stack.append(TestStateResponseReturn(mock))
+            return
+
+        if mock_error_output := mock_input.get("errorOutput"):
+            mock = copy.deepcopy(mock_error_output)
+            self._result_stack.append(TestStateResponseThrow(**mock))
+            return
+
+    def is_mocked(self):
+        if self._mock_input or self._state_configuration:
+            return True
+
+        return False
+
+    def add_result(self, result: TestStateMockedResponse):
+        mock = copy.deepcopy(result)
+        self._result_stack.append(mock)
+
+    def get_next_result(self) -> TestStateMockedResponse:
+        if not self._result_stack:
+            return None
+        return self._result_stack.pop()
+
+    def get_context(self) -> ContextObjectData | None:
+        if self._context is not None:
+            return copy.deepcopy(self._context)
+        return None
+
+    @staticmethod
+    def parse_context(context: str, state_type: StateType = None) -> ContextObjectData:
+        """Parse and validate context JSON string."""
+        try:
+            validation_result = TestStateContextObjectValidator.model_validate_json(context)
+            return validation_result.model_dump(exclude_unset=True, exclude_none=True)
+        except ValidationError as e:
+            error = e.errors()[0]
+            path_str = ".".join(str(x) for x in error["loc"])
+
+            match error:
+                case {"type": "extra_forbidden", "loc": ("Map",)}:
+                    raise ValueError("'Map' field is not supported when mocking a Context object")
+
+                case {"type": "extra_forbidden", "loc": (*_, forbidden_key)}:
+                    raise ValueError(f"Field '{forbidden_key}' is not allowed")
+
+                case {"type": t} if t in ("string_type", "int_type", "dict_type", "model_type"):
+                    expected_map = {
+                        "string_type": "string",
+                        "int_type": "integer",
+                        "dict_type": "object",
+                        "model_type": "object",
+                    }
+                    expected = expected_map.get(t, "valid type")
+                    raise ValueError(f"{path_str} must be a {expected}")
+                case _:
+                    raise ValueError(f"{error['msg']}")

localstack/services/stepfunctions/local_mocking/__init__.py

@@ -0,0 +1,9 @@
+"""
+local_mocking
+-------------
+
+The implementation of the Step Functions Local Mocking
+
+Note that Step Functions Local is different from TestState API mocks.
+TestState API mocking works differently and is implemented separately.
+"""

localstack/services/stepfunctions/{mocking → local_mocking}/mock_config.py

@@ -1,7 +1,7 @@
 import abc
 from typing import Any, Final
 
-from localstack.services.stepfunctions.mocking.mock_config_file import (
+from localstack.services.stepfunctions.local_mocking.mock_config_file import (
     RawMockConfig,
     RawResponseModel,
     RawTestCase,
@@ -9,7 +9,7 @@ from localstack.services.stepfunctions.mocking.mock_config_file import (
 )
 
 
-class MockedResponse(abc.ABC):
+class LocalMockedResponse(abc.ABC):
     range_start: Final[int]
     range_end: Final[int]
 
@@ -28,7 +28,7 @@ class MockedResponse(abc.ABC):
         self.range_end = range_end
 
 
-class MockedResponseReturn(MockedResponse):
+class LocalMockedResponseReturn(LocalMockedResponse):
     payload: Final[Any]
 
     def __init__(self, range_start: int, range_end: int, payload: Any):
@@ -36,7 +36,7 @@ class MockedResponseReturn(MockedResponse):
         self.payload = payload
 
 
-class MockedResponseThrow(MockedResponse):
+class LocalMockedResponseThrow(LocalMockedResponse):
     error: Final[str]
     cause: Final[str]
 
@@ -49,10 +49,13 @@ class MockedResponseThrow(MockedResponse):
 class StateMockedResponses:
     state_name: Final[str]
     mocked_response_name: Final[str]
-    mocked_responses: Final[list[MockedResponse]]
+    mocked_responses: Final[list[LocalMockedResponse]]
 
     def __init__(
-        self, state_name: str, mocked_response_name: str, mocked_responses: list[MockedResponse]
+        self,
+        state_name: str,
+        mocked_response_name: str,
+        mocked_responses: list[LocalMockedResponse],
     ):
         self.state_name = state_name
         self.mocked_response_name = mocked_response_name
@@ -74,7 +77,7 @@ class StateMockedResponses:
             last_range_end = mocked_response.range_end
 
 
-class MockTestCase:
+class LocalMockTestCase:
     state_machine_name: Final[str]
     test_case_name: Final[str]
     state_mocked_responses: Final[dict[str, StateMockedResponses]]
@@ -127,13 +130,15 @@ def _parse_mocked_response_range(string_definition: str) -> tuple[int, int]:
 
 def _mocked_response_from_raw(
     raw_response_model_range: str, raw_response_model: RawResponseModel
-) -> MockedResponse:
+) -> LocalMockedResponse:
     range_start, range_end = _parse_mocked_response_range(raw_response_model_range)
     if raw_response_model.Return:
         payload = raw_response_model.Return.model_dump()
-        return MockedResponseReturn(range_start=range_start, range_end=range_end, payload=payload)
+        return LocalMockedResponseReturn(
+            range_start=range_start, range_end=range_end, payload=payload
+        )
     throw_definition = raw_response_model.Throw
-    return MockedResponseThrow(
+    return LocalMockedResponseThrow(
         range_start=range_start,
         range_end=range_end,
         error=throw_definition.Error,
@@ -143,7 +148,7 @@ def _mocked_response_from_raw(
 
 def _mocked_responses_from_raw(
     mocked_response_name: str, raw_mock_config: RawMockConfig
-) -> list[MockedResponse]:
+) -> list[LocalMockedResponse]:
     raw_response_models: dict[str, RawResponseModel] | None = raw_mock_config.MockedResponses.get(
         mocked_response_name
     )
@@ -151,9 +156,9 @@ def _mocked_responses_from_raw(
         raise RuntimeError(
             f"No definitions for mocked response '{mocked_response_name}' in the mock configuration file."
         )
-    mocked_responses: list[MockedResponse] = []
+    mocked_responses: list[LocalMockedResponse] = []
     for raw_response_model_range, raw_response_model in raw_response_models.items():
-        mocked_response: MockedResponse = _mocked_response_from_raw(
+        mocked_response: LocalMockedResponse = _mocked_response_from_raw(
             raw_response_model_range=raw_response_model_range, raw_response_model=raw_response_model
         )
         mocked_responses.append(mocked_response)
@@ -175,7 +180,7 @@ def _state_mocked_responses_from_raw(
 
 def _mock_test_case_from_raw(
     state_machine_name: str, test_case_name: str, raw_mock_config: RawMockConfig
-) -> MockTestCase:
+) -> LocalMockTestCase:
     state_machine = raw_mock_config.StateMachines.get(state_machine_name)
     if not state_machine:
         raise RuntimeError(
@@ -195,18 +200,20 @@ def _mock_test_case_from_raw(
             raw_mock_config=raw_mock_config,
         )
         state_mocked_responses_list.append(state_mocked_responses)
-    return MockTestCase(
+    return LocalMockTestCase(
         state_machine_name=state_machine_name,
         test_case_name=test_case_name,
         state_mocked_responses_list=state_mocked_responses_list,
     )
 
 
-def load_mock_test_case_for(state_machine_name: str, test_case_name: str) -> MockTestCase | None:
+def load_local_mock_test_case_for(
+    state_machine_name: str, test_case_name: str
+) -> LocalMockTestCase | None:
     raw_mock_config: RawMockConfig | None = _load_sfn_raw_mock_config()
     if raw_mock_config is None:
         return None
-    mock_test_case: MockTestCase = _mock_test_case_from_raw(
+    mock_test_case: LocalMockTestCase = _mock_test_case_from_raw(
         state_machine_name=state_machine_name,
         test_case_name=test_case_name,
         raw_mock_config=raw_mock_config,

localstack/services/stepfunctions/provider.py

@@ -57,13 +57,13 @@ from localstack.aws.api.stepfunctions import (
     LongArn,
     MaxConcurrency,
     MissingRequiredParameter,
+    MockInput,
     Name,
     PageSize,
     PageToken,
     Publish,
     PublishStateMachineVersionOutput,
     ResourceNotFound,
-    RevealSecrets,
     ReverseOrder,
     RevisionId,
     RoutingConfigurationList,
@@ -89,6 +89,7 @@ from localstack.aws.api.stepfunctions import (
     TaskDoesNotExist,
     TaskTimedOut,
     TaskToken,
+    TestStateInput,
     TestStateOutput,
     ToleratedFailureCount,
     ToleratedFailurePercentage,
@@ -150,9 +151,10 @@ from localstack.services.stepfunctions.backend.store import SFNStore, sfn_stores
 from localstack.services.stepfunctions.backend.test_state.execution import (
     TestStateExecution,
 )
-from localstack.services.stepfunctions.mocking.mock_config import (
-    MockTestCase,
-    load_mock_test_case_for,
+from localstack.services.stepfunctions.backend.test_state.test_state_mock import TestStateMock
+from localstack.services.stepfunctions.local_mocking.mock_config import (
+    LocalMockTestCase,
+    load_local_mock_test_case_for,
 )
 from localstack.services.stepfunctions.stepfunctions_utils import (
     assert_pagination_parameters_valid,
@@ -240,6 +242,12 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
             f"Invalid Arn: 'Resource type not valid in this context: {lower_resource_type}'"
         )
 
+    @staticmethod
+    def _validate_test_state_mock_input(mock_input: MockInput) -> None:
+        if {"result", "errorOutput"} <= mock_input.keys():
+            # FIXME create proper error
+            raise ValidationException("Cannot define both 'result' and 'errorOutput'")
+
     @staticmethod
     def _validate_activity_name(name: str) -> None:
         # The activity name is validated according to the AWS StepFunctions documentation, the name should not contain:
@@ -772,14 +780,16 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
         return state_machine_arn.split("#")[0]
 
     @staticmethod
-    def _get_mock_test_case(state_machine_arn: str, state_machine_name: str) -> MockTestCase | None:
+    def _get_local_mock_test_case(
+        state_machine_arn: str, state_machine_name: str
+    ) -> LocalMockTestCase | None:
         """Extract and load a mock test case from a state machine ARN if present."""
         parts = state_machine_arn.split("#")
         if len(parts) != 2:
             return None
 
         mock_test_case_name = parts[1]
-        mock_test_case = load_mock_test_case_for(
+        mock_test_case = load_local_mock_test_case_for(
             state_machine_name=state_machine_name, test_case_name=mock_test_case_name
         )
         if mock_test_case is None:
@@ -856,7 +866,9 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
                 configuration=state_machine_clone.cloud_watch_logging_configuration,
             )
 
-        mock_test_case = self._get_mock_test_case(state_machine_arn, state_machine_clone.name)
+        local_mock_test_case = self._get_local_mock_test_case(
+            state_machine_arn, state_machine_clone.name
+        )
 
         execution = Execution(
             name=exec_name,
@@ -872,7 +884,7 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
             input_data=input_data,
             trace_header=trace_header,
             activity_store=self.get_store(context).activities,
-            mock_test_case=mock_test_case,
+            local_mock_test_case=local_mock_test_case,
         )
 
         store.executions[exec_arn] = execution
@@ -932,7 +944,9 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
                 configuration=state_machine_clone.cloud_watch_logging_configuration,
             )
 
-        mock_test_case = self._get_mock_test_case(state_machine_arn, state_machine_clone.name)
+        local_mock_test_case = self._get_local_mock_test_case(
+            state_machine_arn, state_machine_clone.name
+        )
 
         execution = SyncExecution(
             name=exec_name,
@@ -947,7 +961,7 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
             input_data=input_data,
             trace_header=trace_header,
             activity_store=self.get_store(context).activities,
-            mock_test_case=mock_test_case,
+            local_mock_test_case=local_mock_test_case,
         )
         self.get_store(context).executions[exec_arn] = execution
 
@@ -1483,20 +1497,48 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
         raise ResourceNotFound()
 
     def test_state(
-        self,
-        context: RequestContext,
-        definition: Definition,
-        role_arn: Arn = None,
-        input: SensitiveData = None,
-        inspection_level: InspectionLevel = None,
-        reveal_secrets: RevealSecrets = None,
-        variables: SensitiveData = None,
-        **kwargs,
+        self, context: RequestContext, request: TestStateInput, **kwargs
     ) -> TestStateOutput:
+        state_name = request.get("stateName")
+        definition = request["definition"]
+
         StepFunctionsProvider._validate_definition(
-            definition=definition, static_analysers=[TestStateStaticAnalyser()]
+            definition=definition,
+            static_analysers=[TestStateStaticAnalyser(state_name)],
         )
 
+        # if StateName is present, we need to ensure the state being referenced exists in full definition.
+        if state_name and not TestStateStaticAnalyser.is_state_in_definition(
+            definition=definition, state_name=state_name
+        ):
+            raise ValidationException("State not found in definition")
+
+        mock_input = request.get("mock")
+        if mock_input is not None:
+            self._validate_test_state_mock_input(mock_input)
+            TestStateStaticAnalyser.validate_mock(
+                mock_input=mock_input, definition=definition, state_name=state_name
+            )
+
+        if state_configuration := request.get("stateConfiguration"):
+            # TODO: Add validations for this i.e assert len(input) <= failureCount
+            pass
+
+        if state_context := request.get("context"):
+            # TODO: Add validation ensuring only present if 'mock' is specified
+            # An error occurred (ValidationException) when calling the TestState operation: State type 'Pass' is not supported when a mock is specified
+            pass
+
+        try:
+            state_mock = TestStateMock(
+                mock_input=mock_input,
+                state_configuration=state_configuration,
+                context=state_context,
+            )
+        except ValueError as e:
+            LOG.error(e)
+            raise ValidationException(f"Invalid Context object provided: {e}")
+
         name: Name | None = f"TestState-{short_uid()}"
         arn = stepfunctions_state_machine_arn(
             name=name, account_id=context.account_id, region_name=context.region
@@ -1504,27 +1546,36 @@ class StepFunctionsProvider(StepfunctionsApi, ServiceLifecycleHook):
         state_machine = TestStateMachine(
             name=name,
             arn=arn,
-            role_arn=role_arn,
-            definition=definition,
+            role_arn=request["roleArn"],
+            definition=request["definition"],
         )
-        exec_arn = stepfunctions_standard_execution_arn(state_machine.arn, name)
 
-        input_json = json.loads(input)
+        # HACK(gregfurman): The ARN that gets generated has a duplicate 'name' field in the
+        # resource ARN. Just replace this duplication and extract the execution ID.
+        exec_arn = stepfunctions_express_execution_arn(state_machine.arn, name)
+        exec_arn = exec_arn.replace(f":{name}:{name}:", f":{name}:", 1)
+        _, exec_name = exec_arn.rsplit(":", 1)
+
+        if input_json := request.get("input", {}):
+            input_json = json.loads(input_json)
+
         execution = TestStateExecution(
-            name=name,
-            role_arn=role_arn,
+            name=exec_name,
+            role_arn=request["roleArn"],
             exec_arn=exec_arn,
             account_id=context.account_id,
             region_name=context.region,
             state_machine=state_machine,
             start_date=datetime.datetime.now(tz=datetime.UTC),
             input_data=input_json,
+            state_name=state_name,
             activity_store=self.get_store(context).activities,
+            mock=state_mock,
         )
         execution.start()
 
         test_state_output = execution.to_test_state_output(
-            inspection_level=inspection_level or InspectionLevel.INFO
+            inspection_level=request.get("inspectionLevel", InspectionLevel.INFO)
         )
 
         return test_state_output

localstack/services/stepfunctions/test_state/mock_config.py

@@ -0,0 +1,47 @@
+import abc
+from typing import Any, Final
+
+from pydantic import BaseModel, ConfigDict, StrictInt, StrictStr, create_model
+
+from localstack.services.stepfunctions.asl.eval.states import (
+    ExecutionData,
+    StateData,
+    StateMachineData,
+    TaskData,
+)
+
+
+class TestStateMockedResponse(abc.ABC):
+    pass
+
+
+class TestStateResponseReturn(TestStateMockedResponse):
+    payload: Final[Any]
+
+    def __init__(self, payload: Any):
+        self.payload = payload
+
+
+class TestStateResponseThrow(TestStateMockedResponse):
+    error: Final[str]
+    cause: Final[str]
+
+    def __init__(self, error: str, cause: str):
+        self.error = error
+        self.cause = cause
+
+
+def _to_strict_model(name: str, source: type):
+    type_map = {str: StrictStr, int: StrictInt}
+    fields = {k: (type_map.get(v, v) | None, None) for k, v in source.__annotations__.items()}
+    return create_model(name, __config__=ConfigDict(extra="forbid"), **fields)
+
+
+TestStateContextObjectValidator: Final[type[BaseModel]] = create_model(
+    "ContextValidator",
+    __config__=ConfigDict(extra="forbid"),
+    Execution=(_to_strict_model("Execution", ExecutionData) | None, None),
+    State=(_to_strict_model("State", StateData) | None, None),
+    StateMachine=(_to_strict_model("StateMachine", StateMachineData) | None, None),
+    Task=(_to_strict_model("Task", TaskData) | None, None),
+)