def test_execute_none_should_fail(self, mock_session, mock_dag_bag):
params_list = None
with self.assertRaises(Exception):
operator = MultiTriggerDagRunOperator(
task_id=TASK_ID,
trigger_dag_id=TRIGGER_DAG_ID,
params_list=params_list,
default_args=DAG_ARGS)
operator.execute(None)
mock_session.add.assert_not_called()
mock_session.commit.assert_not_called()
def test_should_execute_iterable(self, mock_session, dag_bag_class):
params_list = xrange(1, 10)
dag_bag_class.return_value = TestMultiTriggerDag.create_mock_dag_bag()
operator = MultiTriggerDagRunOperator(
task_id=TASK_ID,
trigger_dag_id=TRIGGER_DAG_ID,
params_list=params_list,
default_args=DAG_ARGS)
operator.execute(None)
TestMultiTriggerDag.verify_session(params_list)
def test_should_execute_thunked_generator(self, mock_session,
dag_bag_class):
def param_generator():
iterable = range(1, 10)
for i in iterable:
yield i
dag_bag_class.return_value = TestMultiTriggerDag.create_mock_dag_bag()
operator = MultiTriggerDagRunOperator(task_id=TASK_ID,
trigger_dag_id=TRIGGER_DAG_ID,
params_list=param_generator,
default_args=DAG_ARGS)
operator.execute(None)
TestMultiTriggerDag.verify_session(param_generator())
def test_should_add_params_list(self, mock_session, dag_bag_class):
a = "a"
b = "b"
c = "c"
d = "d"
params_list = [a, b, c, d]
dag_bag_class.return_value = TestMultiTriggerDag.create_mock_dag_bag()
operator = MultiTriggerDagRunOperator(task_id=TASK_ID,
trigger_dag_id=TRIGGER_DAG_ID,
params_list=params_list,
default_args=DAG_ARGS)
operator.execute(None)
TestMultiTriggerDag.verify_session(params_list)
# ####################### SCHEDULER #################################
scheduler_dag = DAG(dag_id=SCHEDULE_DAG_ID,
default_args=default_args,
schedule_interval=None)
def get_generator():
iterable = range(0, 100)
for i in iterable:
yield i
operator = MultiTriggerDagRunOperator(task_id='trigger_%s' % TARGET_DAG_ID,
trigger_dag_id=TARGET_DAG_ID,
params_list=get_generator,
default_args=default_args,
dag=scheduler_dag)
# ####################### TARGET DAG #################################
target_dag = DAG(dag_id=TARGET_DAG_ID,
default_args=default_args,
schedule_interval=None)
start = BashOperator(task_id='bash_task',
bash_command='sleep 1; echo "Hello from message #' +
'{{ dag_run.conf if dag_run else "NO MESSAGE" }}"',
default_args=default_args,
dag=target_dag)