def test_trigger_dag_dag_run_exist(self, dag_bag_mock, dag_run_mock):
dag_id = "dag_run_exist"
dag = DAG(dag_id)
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
dag_run_mock.find.return_value = DagRun()
with self.assertRaises(AirflowException):
_trigger_dag(dag_id, dag_bag_mock)
def test_trigger_dag_with_too_early_start_date(self, dag_bag_mock):
dag_id = "trigger_dag_with_too_early_start_date"
dag = DAG(dag_id,
default_args={
'start_date': timezone.datetime(2016, 9, 5, 10, 10, 0)
})
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
with self.assertRaises(ValueError):
_trigger_dag(dag_id,
dag_bag_mock,
execution_date=timezone.datetime(
2015, 7, 5, 10, 10, 0))
def test_trigger_dag_with_conf(self, conf, expected_conf, dag_bag_mock):
dag_id = "trigger_dag_with_conf"
dag = DAG(dag_id)
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
dag_bag_mock.dags_hash = {}
triggers = _trigger_dag(dag_id, dag_bag_mock, conf=conf)
self.assertEqual(triggers[0].conf, expected_conf)
def test_trigger_dag_include_nested_subdags(self, dag_bag_mock,
dag_run_mock, dag_mock):
dag_id = "trigger_dag"
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag_mock
dag_run_mock.find.return_value = None
dag1 = mock.MagicMock(subdags=[])
dag2 = mock.MagicMock(subdags=[dag1])
dag_mock.subdags = [dag1, dag2]
triggers = _trigger_dag(dag_id, dag_bag_mock)
self.assertEqual(3, len(triggers))
def trigger_dag(self, dag_id, run_id, conf):
try:
dag_path = DagModel.get_current(dag_id).fileloc
except Exception:
dag_path = path.join(DAGS_FOLDER, dag_id + ".py")
triggers = trigger_dag._trigger_dag(
dag_id=dag_id,
dag_run=DagRun(),
dag_bag=DagBag(dag_folder=dag_path),
run_id=run_id,
conf=conf,
execution_date=None,
replace_microseconds=False)
return triggers[0] if triggers else None
def test_trigger_dag_with_conf(self, conf, expected_conf, dag_bag_mock):
dag_id = "trigger_dag_with_conf"
dag = DAG(dag_id)
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
dag_run = DagRun()
triggers = _trigger_dag(dag_id,
dag_bag_mock,
dag_run,
run_id=None,
conf=conf,
execution_date=None,
replace_microseconds=True)
self.assertEqual(triggers[0].conf, expected_conf)
def test_trigger_dag_with_str_conf(self, dag_bag_mock):
dag_id = "trigger_dag_with_str_conf"
dag = DAG(dag_id)
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
conf = "{\"foo\": \"bar\"}"
dag_run = DagRun()
triggers = _trigger_dag(dag_id,
dag_bag_mock,
dag_run,
run_id=None,
conf=conf,
execution_date=None,
replace_microseconds=True)
self.assertEqual(triggers[0].conf, json.loads(conf))
def test_trigger_dag_with_valid_start_date(self, dag_bag_mock):
dag_id = "trigger_dag_with_valid_start_date"
dag = DAG(dag_id,
default_args={
'start_date': timezone.datetime(2016, 9, 5, 10, 10, 0)
})
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
dag_bag_mock.dags_hash = {}
triggers = _trigger_dag(dag_id,
dag_bag_mock,
execution_date=timezone.datetime(
2018, 7, 5, 10, 10, 0))
assert len(triggers) == 1
def test_trigger_dag_with_str_conf(self, dag_bag_mock):
dag_id = "trigger_dag_with_str_conf"
dag = DAG(dag_id)
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
conf = "{\"foo\": \"bar\"}"
dag_run = DagRun()
triggers = _trigger_dag(
dag_id,
dag_bag_mock,
dag_run,
run_id=None,
conf=conf,
execution_date=None,
replace_microseconds=True)
self.assertEquals(triggers[0].conf, json.loads(conf))
def test_trigger_dag_with_dict_conf(self, dag_bag_mock):
dag_id = "trigger_dag_with_dict_conf"
dag = DAG(dag_id)
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
conf = dict(foo="bar")
dag_run = DagRun()
triggers = _trigger_dag(
dag_id,
dag_bag_mock,
dag_run,
run_id=None,
conf=conf,
execution_date=None,
replace_microseconds=True)
self.assertEquals(triggers[0].conf, conf)
def test_trigger_dag_with_valid_start_date(self, dag_bag_mock):
dag_id = "trigger_dag_with_valid_start_date"
dag = DAG(dag_id, default_args={'start_date': timezone.datetime(2016, 9, 5, 10, 10, 0)})
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag
dag_run = DagRun()
triggers = _trigger_dag(
dag_id,
dag_bag_mock,
dag_run,
run_id=None,
conf=None,
execution_date=timezone.datetime(2018, 7, 5, 10, 10, 0),
replace_microseconds=True,
)
assert len(triggers) == 1
def test_trigger_dag_include_nested_subdags(self, dag_bag_mock,
dag_run_mock, dag_mock):
dag_id = "trigger_dag"
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag_mock
dag_run_mock.find.return_value = None
dag1 = mock.MagicMock()
dag1.subdags = []
dag2 = mock.MagicMock()
dag2.subdags = [dag1]
dag_mock.subdags = [dag1, dag2]
triggers = _trigger_dag(dag_id,
dag_bag_mock,
dag_run_mock,
run_id=None,
conf=None,
execution_date=None,
replace_microseconds=True)
self.assertEqual(3, len(triggers))
def test_trigger_dag_include_subdags(self, dag_bag_mock, dag_run_mock, dag_mock):
dag_id = "trigger_dag"
dag_bag_mock.dags = [dag_id]
dag_bag_mock.get_dag.return_value = dag_mock
dag_run_mock.find.return_value = None
dag1 = mock.MagicMock()
dag1.subdags = []
dag2 = mock.MagicMock()
dag2.subdags = []
dag_mock.subdags = [dag1, dag2]
triggers = _trigger_dag(
dag_id,
dag_bag_mock,
dag_run_mock,
run_id=None,
conf=None,
execution_date=None,
replace_microseconds=True)
self.assertEqual(3, len(triggers))
def trigger_dag(self, dag_id, run_id, conf):
try:
dag_path = DagModel.get_current(dag_id).fileloc
except Exception:
dag_path = path.join(DAGS_FOLDER, dag_id + ".py")
dag_bag = DagBag(dag_folder=dag_path)
if not dag_bag.dags:
logging.info("Failed to import dag due to the following errors")
logging.info(dag_bag.import_errors)
logging.info("Sleep for 3 seconds and give it a second try")
sleep(3)
dag_bag = DagBag(dag_folder=dag_path)
triggers = trigger_dag._trigger_dag(dag_id=dag_id,
dag_run=DagRun(),
dag_bag=dag_bag,
run_id=run_id,
conf=conf,
execution_date=None,
replace_microseconds=False)
return triggers[0] if triggers else None
def test_trigger_dag_dag_not_found(self, dag_bag_mock):
dag_bag_mock.dags = {}
with self.assertRaises(AirflowException):
_trigger_dag('dag_not_found', dag_bag_mock)
