def test_run_task(self, task_instance_mock): ti_key = "key" job_id = " job_id" task_instance_mock.key = ti_key task_instance_mock.job_id = job_id executor = DebugExecutor() executor.running = {ti_key} succeeded = executor._run_task(task_instance_mock) assert succeeded task_instance_mock._run_raw_task.assert_called_once_with(job_id=job_id)
def test_sync(self, run_task_mock): run_task_mock.return_value = True executor = DebugExecutor() ti1 = MagicMock(key="t1") ti2 = MagicMock(key="t2") executor.tasks_to_run = [ti1, ti2] executor.sync() assert not executor.tasks_to_run run_task_mock.assert_has_calls([mock.call(ti1), mock.call(ti2)])
def test_queue_task_instance(self): key = "ti_key" ti = MagicMock(key=key) executor = DebugExecutor() executor.queue_task_instance(task_instance=ti, mark_success=True, pool="pool") assert key in executor.queued_tasks assert key in executor.tasks_params assert executor.tasks_params[key] == { "mark_success": True, "pool": "pool", }
def dag_test(args, session=None): """Execute one single DagRun for a given DAG and execution date, using the DebugExecutor.""" dag = get_dag(subdir=args.subdir, dag_id=args.dag_id) dag.clear(start_date=args.execution_date, end_date=args.execution_date, dag_run_state=State.NONE) try: dag.run(executor=DebugExecutor(), start_date=args.execution_date, end_date=args.execution_date) except BackfillUnfinished as e: print(str(e)) show_dagrun = args.show_dagrun imgcat = args.imgcat_dagrun filename = args.save_dagrun if show_dagrun or imgcat or filename: tis = session.query(TaskInstance).filter( TaskInstance.dag_id == args.dag_id, TaskInstance.execution_date == args.execution_date, ).all() dot_graph = render_dag(dag, tis=tis) print() if filename: _save_dot_to_file(dot_graph, filename) if imgcat: _display_dot_via_imgcat(dot_graph) if show_dagrun: print(dot_graph.source)
def dag_test(args): """Execute one single DagRun for a given DAG and execution date, using the DebugExecutor.""" dag = get_dag(subdir=args.subdir, dag_id=args.dag_id) dag.clear(start_date=args.execution_date, end_date=args.execution_date, reset_dag_runs=True) dag.run(executor=DebugExecutor(), start_date=args.execution_date, end_date=args.execution_date)
def test_fail_fast(self, change_state_mock): with mock.patch.dict("os.environ", {"AIRFLOW__DEBUG__FAIL_FAST": "True"}): executor = DebugExecutor() ti1 = MagicMock(key="t1") ti2 = MagicMock(key="t2") ti1._run_raw_task.side_effect = Exception executor.tasks_to_run = [ti1, ti2] executor.sync() assert executor.fail_fast assert not executor.tasks_to_run change_state_mock.assert_has_calls([ mock.call(ti1.key, State.FAILED), mock.call(ti2.key, State.UPSTREAM_FAILED), ])
def test_end(self): ti = MagicMock(key="ti_key") executor = DebugExecutor() executor.tasks_to_run = [ti] executor.running = {ti.key} executor.end() ti.set_state.assert_called_once_with(State.UPSTREAM_FAILED) assert not executor.running
def test_dag_sample_w_template_actual_run(): dagbag = models.DagBag(dag_folder=DAG_DIR, include_examples=False) dag = dagbag.get_dag(dag_id="dag_sample_w_template") # type: models.DAG dag.run( start_date=DEFAULT_DATE, ignore_first_depends_on_past=True, verbose=True, executor=DebugExecutor(), ) session = settings.Session() # type: SASession dagruns = session.query(DagRun) \ .filter(DagRun.dag_id == dag.dag_id) \ .order_by(DagRun.execution_date) \ .all() # type: List[models.DagRun] assert len(dagruns) == 1 assert dagruns[0].execution_date == DEFAULT_DATE assert dagruns[0].state == state.State.SUCCESS
def test_trigger_tasks(self): execute_async_mock = MagicMock() executor = DebugExecutor() executor.execute_async = execute_async_mock executor.queued_tasks = { "t1": (None, 1, None, MagicMock(key="t1")), "t2": (None, 2, None, MagicMock(key="t2")), } executor.trigger_tasks(open_slots=4) assert not executor.queued_tasks assert len(executor.running) == 2 assert len(executor.tasks_to_run) == 2 assert not execute_async_mock.called
def _get_executor(executor_name): """ Creates a new instance of the named executor. In case the executor name is not know in airflow, look for it in the plugins """ if executor_name == Executors.LocalExecutor: return LocalExecutor() elif executor_name == Executors.SequentialExecutor: return SequentialExecutor() elif executor_name == Executors.CeleryExecutor: from airflow.executors.celery_executor import CeleryExecutor return CeleryExecutor() elif executor_name == Executors.DaskExecutor: from airflow.executors.dask_executor import DaskExecutor return DaskExecutor() elif executor_name == Executors.MesosExecutor: from airflow.contrib.executors.mesos_executor import MesosExecutor return MesosExecutor() elif executor_name == Executors.KubernetesExecutor: from airflow.contrib.executors.kubernetes_executor import KubernetesExecutor return KubernetesExecutor() elif executor_name == Executors.DebugExecutor: from airflow.executors.debug_executor import DebugExecutor return DebugExecutor() else: # Loading plugins _integrate_plugins() executor_path = executor_name.split('.') if len(executor_path) != 2: raise AirflowException( "Executor {0} not supported: " "please specify in format plugin_module.executor".format( executor_name)) if executor_path[0] in globals(): return globals()[executor_path[0]].__dict__[executor_path[1]]() else: raise AirflowException( "Executor {0} not supported.".format(executor_name))
def dag_test(args, session=None): """Execute one single DagRun for a given DAG and execution date, using the DebugExecutor.""" dag = get_dag(subdir=args.subdir, dag_id=args.dag_id) dag.clear(start_date=args.execution_date, end_date=args.execution_date, dag_run_state=False) try: dag.run( executor=DebugExecutor(), start_date=args.execution_date, end_date=args.execution_date, # Always run the DAG at least once even if no logical runs are # available. This does not make a lot of sense, but Airflow has # been doing this prior to 2.2 so we keep compatibility. run_at_least_once=True, ) except BackfillUnfinished as e: print(str(e)) show_dagrun = args.show_dagrun imgcat = args.imgcat_dagrun filename = args.save_dagrun if show_dagrun or imgcat or filename: tis = ( session.query(TaskInstance) .filter( TaskInstance.dag_id == args.dag_id, TaskInstance.execution_date == args.execution_date, ) .all() ) dot_graph = render_dag(dag, tis=tis) print() if filename: _save_dot_to_file(dot_graph, filename) if imgcat: _display_dot_via_imgcat(dot_graph) if show_dagrun: print(dot_graph.source)