Skip to content
This repository was archived by the owner on Mar 23, 2026. It is now read-only.

Commit c92c8e2

Browse files
tiurintiurin
authored
SFN:TestState: Add mock result validations for Map state (#13492)
Adds validations conditioned by the presence of ResultWriter field. Refactors existing validations to reuse mock parsing and reduce nested `if` blocks.
1 parent 16356f0 commit c92c8e2

File tree

6 files changed

+279
-142
lines changed

6 files changed

+279
-142
lines changed

‎localstack-core/localstack/services/stepfunctions/asl/static_analyser/test_state/test_state_analyser.py‎

Lines changed: 128 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,28 @@ def validate_mock(test_state_input: TestStateInput) -> None:
109109
"A test mock should have only one of the following fields: [result, errorOutput]."
110110
)
111111

112-
TestStateStaticAnalyser.validate_mock_result_matches_api_shape(
113-
mock_input=mock_input, test_state=test_state
114-
)
112+
mock_result_raw = mock_input.get("result")
113+
if mock_result_raw is None:
114+
return
115+
try:
116+
mock_result = json.loads(mock_result_raw)
117+
except json.JSONDecodeError:
118+
raise ValidationException("Mocked result must be valid JSON")
119+
120+
if isinstance(test_state, StateMap):
121+
TestStateStaticAnalyser.validate_mock_result_matches_map_definition(
122+
mock_result=mock_result, test_state=test_state
123+
)
124+
125+
if isinstance(test_state, StateTaskService):
126+
field_validation_mode = mock_input.get(
127+
"fieldValidationMode", MockResponseValidationMode.STRICT
128+
)
129+
TestStateStaticAnalyser.validate_mock_result_matches_api_shape(
130+
mock_result=mock_result,
131+
field_validation_mode=field_validation_mode,
132+
test_state=test_state,
133+
)
115134

116135
@staticmethod
117136
def validate_test_state_allows_mocking(
@@ -130,127 +149,126 @@ def validate_test_state_allows_mocking(
130149
)
131150

132151
@staticmethod
133-
def validate_mock_result_matches_api_shape(mock_input: MockInput, test_state: CommonStateField):
152+
def validate_mock_result_matches_map_definition(mock_result: Any, test_state: StateMap):
153+
if test_state.result_writer is not None and not isinstance(mock_result, dict):
154+
raise ValidationException("Mocked result must be a JSON object.")
155+
156+
if test_state.result_writer is None and not isinstance(mock_result, list):
157+
raise ValidationException("Mocked result must be an array.")
158+
159+
@staticmethod
160+
def validate_mock_result_matches_api_shape(
161+
mock_result: Any,
162+
field_validation_mode: MockResponseValidationMode,
163+
test_state: StateTaskService,
164+
):
134165
# apigateway:invoke: has no equivalent in the AWS SDK service integration.
135166
# Hence, the validation against botocore doesn't apply.
136167
# See the note in https://docs.aws.amazon.com/step-functions/latest/dg/connect-api-gateway.html
137168
# TODO do custom validation for apigateway:invoke:
138169
if isinstance(test_state, StateTaskServiceApiGateway):
139170
return
140171

141-
if isinstance(test_state, StateTaskService):
142-
field_validation_mode = mock_input.get(
143-
"fieldValidationMode", MockResponseValidationMode.STRICT
172+
if field_validation_mode == MockResponseValidationMode.NONE:
173+
return
174+
175+
boto_service_name = test_state._get_boto_service_name()
176+
service_action_name = test_state._get_boto_service_action()
177+
output_shape = test_state._get_boto_operation_model(
178+
boto_service_name=boto_service_name, service_action_name=service_action_name
179+
).output_shape
180+
181+
# If the operation has no output, there's nothing to validate
182+
if output_shape is None:
183+
return
184+
185+
def _raise_type_error(expected_type: str, field_name: str) -> None:
186+
raise ValidationException(
187+
f"Mock result schema validation error: Field '{field_name}' must be {expected_type}"
144188
)
145-
mock_result_raw = mock_input.get("result")
146-
if mock_result_raw is None:
189+
190+
def _validate_value(value: Any, shape: Shape, field_name: str | None = None) -> None:
191+
# Document type accepts any JSON value
192+
if shape.type_name == "document":
147193
return
148-
try:
149-
mock_result = json.loads(mock_result_raw)
150-
except json.JSONDecodeError:
151-
raise ValidationException("Mocked result must be valid JSON")
152-
if mock_result is None or field_validation_mode == MockResponseValidationMode.NONE:
194+
195+
if isinstance(shape, StructureShape):
196+
if not isinstance(value, dict):
197+
# this is a defensive check, the mock result is loaded from JSON before, so should always be a dict
198+
raise ValidationException(
199+
f"Mock result must be a valid JSON object, but got '{type(value)}' instead"
200+
)
201+
# Build a mapping from SFN-normalised member keys -> botocore member shapes
202+
members = shape.members
203+
sfn_key_to_member_shape: dict[str, Shape] = {
204+
StateTaskService._to_sfn_cased(member_key): member_shape
205+
for member_key, member_shape in members.items()
206+
}
207+
if field_validation_mode == MockResponseValidationMode.STRICT:
208+
# Ensure required members are present, using SFN-normalised keys
209+
for required_key in shape.required_members:
210+
sfn_required_key = StateTaskService._to_sfn_cased(required_key)
211+
if sfn_required_key not in value:
212+
raise ValidationException(
213+
f"Mock result schema validation error: Required field '{sfn_required_key}' is missing"
214+
)
215+
# Validate present fields (match SFN-normalised keys to member shapes)
216+
for mock_field_name, mock_field_value in value.items():
217+
member_shape = sfn_key_to_member_shape.get(mock_field_name)
218+
if member_shape is None:
219+
# Fields that are present in mock but are not in the API spec should not raise validation errors - forward compatibility
220+
continue
221+
_validate_value(mock_field_value, member_shape, mock_field_name)
153222
return
154223

155-
boto_service_name = test_state._get_boto_service_name()
156-
service_action_name = test_state._get_boto_service_action()
157-
output_shape = test_state._get_boto_operation_model(
158-
boto_service_name=boto_service_name, service_action_name=service_action_name
159-
).output_shape
224+
if isinstance(shape, ListShape):
225+
if not isinstance(value, list):
226+
_raise_type_error("an array", field_name)
227+
member_shape = shape.member
228+
for list_item in value:
229+
_validate_value(list_item, member_shape, field_name)
230+
return
160231

161-
# If the operation has no output, there's nothing to validate
162-
if output_shape is None:
232+
if isinstance(shape, MapShape):
233+
if not isinstance(value, dict):
234+
_raise_type_error("an object", field_name)
235+
value_shape = shape.value
236+
for _, map_item_value in value.items():
237+
_validate_value(map_item_value, value_shape, field_name)
163238
return
164239

165-
def _raise_type_error(expected_type: str, field_name: str) -> None:
166-
raise ValidationException(
167-
f"Mock result schema validation error: Field '{field_name}' must be {expected_type}"
168-
)
169-
170-
def _validate_value(value: Any, shape: Shape, field_name: str | None = None) -> None:
171-
# Document type accepts any JSON value
172-
if shape.type_name == "document":
173-
return
174-
175-
if isinstance(shape, StructureShape):
176-
if not isinstance(value, dict):
177-
# this is a defensive check, the mock result is loaded from JSON before, so should always be a dict
178-
raise ValidationException(
179-
f"Mock result must be a valid JSON object, but got '{type(value)}' instead"
180-
)
181-
# Build a mapping from SFN-normalised member keys -> botocore member shapes
182-
members = shape.members
183-
sfn_key_to_member_shape: dict[str, Shape] = {
184-
StateTaskService._to_sfn_cased(member_key): member_shape
185-
for member_key, member_shape in members.items()
186-
}
187-
if field_validation_mode == MockResponseValidationMode.STRICT:
188-
# Ensure required members are present, using SFN-normalised keys
189-
for required_key in shape.required_members:
190-
sfn_required_key = StateTaskService._to_sfn_cased(required_key)
191-
if sfn_required_key not in value:
192-
raise ValidationException(
193-
f"Mock result schema validation error: Required field '{sfn_required_key}' is missing"
194-
)
195-
# Validate present fields (match SFN-normalised keys to member shapes)
196-
for mock_field_name, mock_field_value in value.items():
197-
member_shape = sfn_key_to_member_shape.get(mock_field_name)
198-
if member_shape is None:
199-
# Fields that are present in mock but are not in the API spec should not raise validation errors - forward compatibility
200-
continue
201-
_validate_value(mock_field_value, member_shape, mock_field_name)
202-
return
203-
204-
if isinstance(shape, ListShape):
205-
if not isinstance(value, list):
206-
_raise_type_error("an array", field_name)
207-
member_shape = shape.member
208-
for list_item in value:
209-
_validate_value(list_item, member_shape, field_name)
210-
return
211-
212-
if isinstance(shape, MapShape):
213-
if not isinstance(value, dict):
214-
_raise_type_error("an object", field_name)
215-
value_shape = shape.value
216-
for _, map_item_value in value.items():
217-
_validate_value(map_item_value, value_shape, field_name)
218-
return
219-
220-
# Primitive shapes and others
221-
type_name = shape.type_name
222-
match type_name:
223-
case "string" | "timestamp":
224-
if not isinstance(value, str):
225-
_raise_type_error("a string", field_name)
226-
# Validate enum if present
227-
if isinstance(shape, StringShape):
228-
enum = getattr(shape, "enum", None)
229-
if enum and value not in enum:
230-
raise ValidationException(
231-
f"Mock result schema validation error: Field '{field_name}' is not an expected value"
232-
)
233-
234-
case "integer" | "long":
235-
if not isinstance(value, int) or isinstance(value, bool):
236-
_raise_type_error("a number", field_name)
237-
238-
case "float" | "double":
239-
if not (isinstance(value, (int, float)) or isinstance(value, bool)):
240-
_raise_type_error("a number", field_name)
241-
242-
case "boolean":
243-
if not isinstance(value, bool):
244-
_raise_type_error("a boolean", field_name)
245-
246-
case "blob":
247-
if not (isinstance(value, (str, bytes))):
248-
_raise_type_error("a string", field_name)
249-
250-
# Perform validation against the output shape
251-
_validate_value(mock_result, output_shape)
252-
# Non-service tasks or other cases: nothing to validate
253-
return
240+
# Primitive shapes and others
241+
type_name = shape.type_name
242+
match type_name:
243+
case "string" | "timestamp":
244+
if not isinstance(value, str):
245+
_raise_type_error("a string", field_name)
246+
# Validate enum if present
247+
if isinstance(shape, StringShape):
248+
enum = getattr(shape, "enum", None)
249+
if enum and value not in enum:
250+
raise ValidationException(
251+
f"Mock result schema validation error: Field '{field_name}' is not an expected value"
252+
)
253+
254+
case "integer" | "long":
255+
if not isinstance(value, int) or isinstance(value, bool):
256+
_raise_type_error("a number", field_name)
257+
258+
case "float" | "double":
259+
if not (isinstance(value, (int, float)) or isinstance(value, bool)):
260+
_raise_type_error("a number", field_name)
261+
262+
case "boolean":
263+
if not isinstance(value, bool):
264+
_raise_type_error("a boolean", field_name)
265+
266+
case "blob":
267+
if not (isinstance(value, (str, bytes))):
268+
_raise_type_error("a string", field_name)
269+
270+
# Perform validation against the output shape
271+
_validate_value(mock_result, output_shape)
254272

255273
def analyse(self, definition: str) -> None:
256274
_, parser_rule_context = TestStateAmazonStateLanguageParser.parse(

‎tests/aws/services/stepfunctions/templates/test_state/statemachines/base_map_state_result_writer.json5‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
{
2+
"Comment": "BASE_MAP_STATE",
3+
"Type": "Map",
4+
"ItemsPath": "$.Values",
5+
"ItemProcessor": {
6+
"ProcessorConfig": {
7+
"Mode": "DISTRIBUTED",
8+
"ExecutionType": "STANDARD"
9+
},
10+
"StartAt": "TestState",
11+
"States": {
12+
"TestState": {
13+
"Type": "Task",
14+
"Resource": "arn:aws:states:::lambda:invoke",
15+
"Parameters": {
16+
"FunctionName": "foo",
17+
"Payload": "bar"
18+
},
19+
"End": true
20+
}
21+
}
22+
},
23+
"ResultWriter": {
24+
"Resource": "arn:aws:states:::s3:putObject",
25+
"Parameters": {
26+
"Bucket": "result-bucket",
27+
"Prefix": "mapJobs"
28+
}
29+
},
30+
"Label": "TestMap",
31+
"End": true
32+
}

‎tests/aws/services/stepfunctions/templates/test_state/test_state_templates.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ class TestStateTemplate(TemplateLoader):
3737
)
3838

3939
BASE_MAP_STATE: Final[str] = os.path.join(_THIS_FOLDER, "statemachines/base_map_state.json5")
40+
BASE_MAP_STATE_WITH_RESULT_WRITER: Final[str] = os.path.join(
41+
_THIS_FOLDER, "statemachines/base_map_state_result_writer.json5"
42+
)
4043
BASE_MAP_STATE_CATCH: Final[str] = os.path.join(
4144
_THIS_FOLDER, "statemachines/base_map_state_catch.json5"
4245
)

‎tests/aws/services/stepfunctions/v2/test_state/test_state_mock_validation.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,37 @@ def test_reveal_secrets_is_true_and_mock_result_is_set(
103103
revealSecrets=True,
104104
)
105105
sfn_snapshot.match("validation_exception", e.value.response)
106+
107+
@markers.aws.validated
108+
def test_mock_result_is_not_array_on_map_state_without_result_writer(
109+
self,
110+
aws_client_no_sync_prefix,
111+
sfn_snapshot,
112+
):
113+
template = TST.load_sfn_template(TST.BASE_MAP_STATE)
114+
definition = json.dumps(template)
115+
mock = {"result": json.dumps({"mock": "array is expected but object is provided instead"})}
116+
with pytest.raises(Exception) as e:
117+
aws_client_no_sync_prefix.stepfunctions.test_state(
118+
definition=definition,
119+
inspectionLevel=InspectionLevel.TRACE,
120+
mock=mock,
121+
)
122+
sfn_snapshot.match("validation_exception", e.value.response)
123+
124+
@markers.aws.validated
125+
def test_mock_result_is_not_object_on_map_state_with_result_writer(
126+
self,
127+
aws_client_no_sync_prefix,
128+
sfn_snapshot,
129+
):
130+
template = TST.load_sfn_template(TST.BASE_MAP_STATE_WITH_RESULT_WRITER)
131+
definition = json.dumps(template)
132+
mock = {"result": json.dumps(["object is expected but array is provided instead"])}
133+
with pytest.raises(Exception) as e:
134+
aws_client_no_sync_prefix.stepfunctions.test_state(
135+
definition=definition,
136+
inspectionLevel=InspectionLevel.TRACE,
137+
mock=mock,
138+
)
139+
sfn_snapshot.match("validation_exception", e.value.response)

0 commit comments

Comments
 (0)