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"}),