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)