def test_command_responds_to_env_name(self):
t1 = AirflowTask(task_id="test-task", dag_id="blob")
assert t1.helper_script == "source deactivate && source activate airflow"
t2 = AirflowTask(task_id="test-task",
dag_id="blob",
airflow_env="airflow_conda_env")
assert (t2.helper_script ==
"source deactivate && source activate airflow_conda_env")
def test_init_requires_task_and_dag_id(self):
with pytest.raises(TypeError):
task = AirflowTask()
with pytest.raises(TypeError):
task = AirflowTask(task_id="task name")
with pytest.raises(TypeError):
task = AirflowTask(dag_id="dag name")
def test_airflow_task_converts_xcoms_to_return_values(
self, airflow_settings):
puller = AirflowTask(
db_conn=airflow_settings["db_conn"],
task_id="puller",
dag_id="example_xcom",
env=airflow_settings,
execution_date="1999-09-20",
)
push = AirflowTask(
db_conn=airflow_settings["db_conn"],
task_id="push",
dag_id="example_xcom",
env=airflow_settings,
execution_date="1999-09-20",
)
push_by_returning = AirflowTask(
db_conn=airflow_settings["db_conn"],
task_id="push_by_returning",
dag_id="example_xcom",
env=airflow_settings,
execution_date="1999-09-20",
)
with Flow(name="xcom") as flow:
res = puller(upstream_tasks=[push, push_by_returning])
flow_state = flow.run()
assert flow_state.is_successful()
# puller
assert flow_state.result[res].is_successful()
assert flow_state.result[res].result is None
# push
assert flow_state.result[push].is_successful()
assert flow_state.result[push].result == [1, 2, 3]
# push_by_returning
assert flow_state.result[push_by_returning].is_successful()
assert flow_state.result[push_by_returning].result == {"a": "b"}
def test_airflow_task_successfully_runs_a_task(self, airflow_settings):
task = AirflowTask(
db_conn=airflow_settings["db_conn"],
task_id="also_run_this",
dag_id="example_bash_operator",
env=airflow_settings,
)
flow = Flow(name="test single task", tasks=[task])
flow_state = flow.run()
assert flow_state.is_successful()
assert flow_state.result[task].is_successful()
assert flow_state.result[task].result is None
def test_airflow_task_uses_its_own_trigger_rules_by_default(
self, airflow_settings):
task = AirflowTask(
db_conn=airflow_settings["db_conn"],
task_id="run_this_last",
dag_id="example_bash_operator",
env=airflow_settings,
)
with Flow(name="test single task") as flow:
res = task(execution_date="2011-01-01")
flow_state = flow.run()
assert flow_state.is_successful()
assert flow_state.result[res].is_skipped()
def test_airflow_task_uses_cli_flags(self, airflow_settings):
task = AirflowTask(
db_conn=airflow_settings["db_conn"],
task_id="run_this_last",
dag_id="example_bash_operator",
cli_flags=["-A"],
env=airflow_settings,
)
with Flow(name="test single task") as flow:
res = task(execution_date="2011-01-02")
flow_state = flow.run()
assert flow_state.is_successful()
assert flow_state.result[res].is_successful()
assert not flow_state.result[res].is_skipped()
assert flow_state.result[res].result is None
def test_command_responds_to_cli_flags(self):
t1 = AirflowTask(task_id="test-task",
dag_id="blob",
cli_flags=["--force"])
assert t1.command.startswith("airflow run --force")
def test_name_defaults_to_task_id_but_can_be_changed(self):
t1 = AirflowTask(task_id="test-task", dag_id="blob")
assert t1.name == "test-task"
t2 = AirflowTask(task_id="test-task", dag_id="blob", name="unique")
assert t2.name == "unique"