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)