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)