def test_map_state_with_placeholders():
workflow_input = ExecutionInput()
step_result = StepResult()
map_state = Map(state_id="MapState01",
result_selector={
"foo": step_result["foo"],
"bar": step_result["bar1"]["bar2"]
})
iterator_state = Pass("TrainIterator",
parameters={
"ParamA": map_state.output()["X"]["Y"],
"ParamB":
workflow_input["Key01"]["Key02"]["Key03"]
})
map_state.attach_iterator(iterator_state)
workflow_definition = Chain([map_state])
expected_repr = {
"StartAt": "MapState01",
"States": {
"MapState01": {
"Type": "Map",
"ResultSelector": {
"foo.$": "$['foo']",
"bar.$": "$['bar1']['bar2']"
},
"End": True,
"Iterator": {
"StartAt": "TrainIterator",
"States": {
"TrainIterator": {
"Parameters": {
"ParamA.$":
"$['X']['Y']",
"ParamB.$":
"$$.Execution.Input['Key01']['Key02']['Key03']"
},
"Type": "Pass",
"End": True
}
}
}
}
}
}
result = Graph(workflow_definition).to_dict()
assert result == expected_repr
def test_map_state_creation():
map_state = Map('Map', iterator=Pass('FirstIteratorState'), items_path='$', max_concurrency=0)
assert map_state.to_dict() == {
'Type': 'Map',
'ItemsPath': '$',
'Iterator': {
'StartAt': 'FirstIteratorState',
'States': {
'FirstIteratorState': {
'Type': 'Pass',
'End': True
}
}
},
'MaxConcurrency': 0,
'End': True
}
def test_map_state_with_placeholders():
workflow_input = ExecutionInput()
map_state = Map('MapState01')
iterator_state = Pass('TrainIterator',
parameters={
'ParamA': map_state.output()['X']["Y"],
'ParamB':
workflow_input["Key01"]["Key02"]["Key03"]
})
map_state.attach_iterator(iterator_state)
workflow_definition = Chain([map_state])
expected_repr = {
"StartAt": "MapState01",
"States": {
"MapState01": {
"Type": "Map",
"End": True,
"Iterator": {
"StartAt": "TrainIterator",
"States": {
"TrainIterator": {
"Parameters": {
"ParamA.$":
"$['X']['Y']",
"ParamB.$":
"$$.Execution.Input['Key01']['Key02']['Key03']"
},
"Type": "Pass",
"End": True
}
}
}
}
}
}
result = Graph(workflow_definition).to_dict()
assert result == expected_repr
def test_map_state_add_catch_adds_catcher_to_catchers(catch, expected_catch):
step = Map('Map', iterator=Pass('Iterator'))
step.add_catch(catch)
assert step.to_dict()['Catch'] == expected_catch
def test_map_state_constructor_with_catch_adds_catcher_to_catchers(
catch, expected_catch):
step = Map('Map', catch=catch, iterator=Pass('Iterator'))
assert step.to_dict()['Catch'] == expected_catch
def test_map_state_add_retry_adds_retrier_to_retriers(retry, expected_retry):
step = Map('Map', iterator=Pass('Iterator'))
step.add_retry(retry)
assert step.to_dict()['Retry'] == expected_retry
def test_map_state_constructor_with_retry_adds_retrier_to_retriers(
retry, expected_retry):
step = Map('Map', retry=retry, iterator=Pass('Iterator'))
assert step.to_dict()['Retry'] == expected_retry