class TestRunFlowStep:
def test_running_state_finishes(self):
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).get_flow_run_state(
state=Running(),
task_states={},
task_contexts={},
return_tasks=set(),
task_runner_state_handlers=[],
executor=LocalExecutor(),
)
assert new_state.is_successful()
@pytest.mark.parametrize(
"state", [Pending(), Retrying(), Finished(), Success(), Failed(), Skipped()]
)
def test_other_states_raise_endrun(self, state):
flow = Flow(name="test", tasks=[Task()])
with pytest.raises(ENDRUN):
FlowRunner(flow=flow).get_flow_run_state(
state=state,
task_states={},
task_contexts={},
return_tasks=set(),
task_runner_state_handlers=[],
executor=Executor(),
)
def test_determine_final_state_has_final_say(self):
class MyFlowRunner(FlowRunner):
def determine_final_state(self, *args, **kwargs):
return Failed("Very specific error message")
flow = Flow(name="test", tasks=[Task()])
new_state = MyFlowRunner(flow=flow).get_flow_run_state(
state=Running(),
task_states={},
task_contexts={},
return_tasks=set(),
task_runner_state_handlers=[],
executor=LocalExecutor(),
)
assert new_state.is_failed()
assert new_state.message == "Very specific error message"
def test_determine_final_state_preserves_running_states_when_tasks_still_running(
self,
):
task = Task()
flow = Flow(name="test", tasks=[task])
old_state = Running()
new_state = FlowRunner(flow=flow).get_flow_run_state(
state=old_state,
task_states={task: Retrying(start_time=pendulum.now("utc").add(days=1))},
task_contexts={},
return_tasks=set(),
task_runner_state_handlers=[],
executor=LocalExecutor(),
)
assert new_state is old_state
def test_graceful_exceptions_handle_failed(mocker, temp_url):
graceful_fail_task_method = mocker.patch('iguazu.core.tasks.Task._graceful_fail',
side_effect=[ENDRUN(state=Finished())])
task = TaskWithException(graceful_exceptions=(CustomException, ))
with Flow('test_graceful_exceptions_handle_failed') as flow:
task(value=+1)
with raise_on_exception(), prefect.context(caches={}), pytest.raises(ValueError):
flow.run()
graceful_fail_task_method.assert_not_called()
class TestCheckFlowPendingOrRunning:
@pytest.mark.parametrize("state", [Pending(), Running(), Retrying(), Scheduled()])
def test_pending_or_running_are_ok(self, state):
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).check_flow_is_pending_or_running(state=state)
assert new_state is state
@pytest.mark.parametrize("state", [Finished(), Success(), Failed(), Skipped()])
def test_not_pending_or_running_raise_endrun(self, state):
flow = Flow(name="test", tasks=[Task()])
with pytest.raises(ENDRUN):
FlowRunner(flow=flow).check_flow_is_pending_or_running(state=state)
def test_graceful_exceptions_handle_graceful(mocker, temp_url):
graceful_fail_task_method = mocker.patch('iguazu.core.tasks.Task._graceful_fail',
side_effect=[ENDRUN(state=Finished())])
task = TaskWithException(graceful_exceptions=(CustomException, ))
with Flow('test_graceful_exceptions_handle_graceful') as flow:
task(value=-1)
with prefect.context(caches={}):
flow.run()
graceful_fail_task_method.assert_called_once()
call_args = graceful_fail_task_method.call_args[0]
assert isinstance(call_args[0], CustomException)
async def test_task_run_is_dequeued_when_state_changes(
self, tenant_id, running_flow_run_id, task_run_id
):
# update the run's state to scheduled
await api.states.set_task_run_state(task_run_id=task_run_id, state=Scheduled())
# assert it is in the queue
runs_in_queue = await api.runs.get_runs_in_queue(tenant_id=tenant_id)
assert running_flow_run_id in runs_in_queue
# update the run's state to finished
await api.states.set_task_run_state(task_run_id=task_run_id, state=Finished())
# assert it is not in the queue
runs_in_queue = await api.runs.get_runs_in_queue(tenant_id=tenant_id)
assert running_flow_run_id not in runs_in_queue
class TestSetFlowToRunning:
@pytest.mark.parametrize("state", [Pending(), Retrying()])
def test_pending_becomes_running(self, state):
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).set_flow_to_running(state=state)
assert new_state.is_running()
def test_running_stays_running(self):
state = Running()
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).set_flow_to_running(state=state)
assert new_state.is_running()
@pytest.mark.parametrize("state", [Finished(), Success(), Failed(), Skipped()])
def test_other_states_raise_endrun(self, state):
flow = Flow(name="test", tasks=[Task()])
with pytest.raises(ENDRUN):
FlowRunner(flow=flow).set_flow_to_running(state=state)
def test_flow_run_accepts_state_kwarg():
f = Flow(name="test")
state = f.run(state=Finished())
assert state.is_finished()
def test_timedout_is_failed(self):
assert issubclass(TimedOut, Failed)
def test_trigger_failed_is_failed(self):
assert issubclass(TriggerFailed, Failed)
@pytest.mark.parametrize(
"state_check",
[
dict(state=Cancelled(), assert_true={"is_finished"}),
dict(state=Cached(),
assert_true={"is_cached", "is_finished", "is_successful"}),
dict(state=ClientFailed(), assert_true={"is_meta_state"}),
dict(state=Failed(), assert_true={"is_finished", "is_failed"}),
dict(state=Finished(), assert_true={"is_finished"}),
dict(state=Looped(), assert_true={"is_finished", "is_looped"}),
dict(state=Mapped(),
assert_true={"is_finished", "is_mapped", "is_successful"}),
dict(state=Paused(), assert_true={"is_pending", "is_scheduled"}),
dict(state=Pending(), assert_true={"is_pending"}),
dict(state=Queued(), assert_true={"is_meta_state", "is_queued"}),
dict(state=Resume(), assert_true={"is_pending", "is_scheduled"}),
dict(state=Retrying(),
assert_true={"is_pending", "is_scheduled", "is_retrying"}),
dict(state=Running(), assert_true={"is_running"}),
dict(state=Scheduled(), assert_true={"is_pending", "is_scheduled"}),
dict(state=Skipped(),
assert_true={"is_finished", "is_successful", "is_skipped"}),
dict(state=Submitted(), assert_true={"is_meta_state", "is_submitted"}),
dict(state=Success(), assert_true={"is_finished", "is_successful"}),