def test_sync_to_db(self, mock_now): dag = DAG( 'dag', start_date=DEFAULT_DATE, ) with dag: DummyOperator(task_id='task', owner='owner1') SubDagOperator( task_id='subtask', owner='owner2', subdag=DAG( 'dag.subtask', start_date=DEFAULT_DATE, ) ) now = datetime.datetime.utcnow().replace(tzinfo=pendulum.timezone('UTC')) mock_now.return_value = now session = settings.Session() dag.sync_to_db(session=session) orm_dag = session.query(DagModel).filter(DagModel.dag_id == 'dag').one() self.assertEqual(set(orm_dag.owners.split(', ')), {'owner1', 'owner2'}) self.assertEqual(orm_dag.last_scheduler_run, now) self.assertTrue(orm_dag.is_active) self.assertIsNone(orm_dag.default_view) self.assertEqual(orm_dag.get_default_view(), configuration.conf.get('webserver', 'dag_default_view').lower()) self.assertEqual(orm_dag.safe_dag_id, 'dag') orm_subdag = session.query(DagModel).filter( DagModel.dag_id == 'dag.subtask').one() self.assertEqual(set(orm_subdag.owners.split(', ')), {'owner1', 'owner2'}) self.assertEqual(orm_subdag.last_scheduler_run, now) self.assertTrue(orm_subdag.is_active) self.assertEqual(orm_subdag.safe_dag_id, 'dag__dot__subtask')
def test_sync_to_db(self, mock_now): dag = DAG( 'dag', start_date=DEFAULT_DATE, ) with dag: DummyOperator(task_id='task', owner='owner1') SubDagOperator(task_id='subtask', owner='owner2', subdag=DAG( 'dag.subtask', start_date=DEFAULT_DATE, )) now = datetime.datetime.utcnow().replace( tzinfo=pendulum.timezone('UTC')) mock_now.return_value = now session = settings.Session() dag.sync_to_db(session=session) orm_dag = session.query(DagModel).filter( DagModel.dag_id == 'dag').one() self.assertEqual(set(orm_dag.owners.split(', ')), {'owner1', 'owner2'}) self.assertEqual(orm_dag.last_scheduler_run, now) self.assertTrue(orm_dag.is_active) self.assertIsNone(orm_dag.default_view) self.assertEqual( orm_dag.get_default_view(), configuration.conf.get('webserver', 'dag_default_view').lower()) self.assertEqual(orm_dag.safe_dag_id, 'dag') orm_subdag = session.query(DagModel).filter( DagModel.dag_id == 'dag.subtask').one() self.assertEqual(set(orm_subdag.owners.split(', ')), {'owner1', 'owner2'}) self.assertEqual(orm_subdag.last_scheduler_run, now) self.assertTrue(orm_subdag.is_active) self.assertEqual(orm_subdag.safe_dag_id, 'dag__dot__subtask') self.assertEqual(orm_subdag.fileloc, orm_dag.fileloc)
def test_dag_is_deactivated_upon_dagfile_deletion(self): dag_id = 'old_existing_dag' dag_fileloc = "/usr/local/airflow/dags/non_existing_path.py" dag = DAG( dag_id, is_paused_upon_creation=True, ) dag.fileloc = dag_fileloc session = settings.Session() dag.sync_to_db(session=session) orm_dag = session.query(DagModel).filter(DagModel.dag_id == dag_id).one() self.assertTrue(orm_dag.is_active) self.assertEqual(orm_dag.fileloc, dag_fileloc) DagModel.deactivate_deleted_dags(list_py_file_paths(settings.DAGS_FOLDER)) orm_dag = session.query(DagModel).filter(DagModel.dag_id == dag_id).one() self.assertFalse(orm_dag.is_active) # CleanUp session.execute(DagModel.__table__.delete().where(DagModel.dag_id == dag_id))
def test_is_paused_subdag(self, mock_dag_bag): subdag_id = 'dag.subdag' subdag = DAG( subdag_id, start_date=DEFAULT_DATE, ) with subdag: DummyOperator(task_id='dummy_task', ) dag_id = 'dag' dag = DAG( dag_id, start_date=DEFAULT_DATE, ) with dag: SubDagOperator(task_id='subdag', subdag=subdag) mock_dag_bag.return_value.get_dag.return_value = dag session = settings.Session() dag.sync_to_db(session=session) unpaused_dags = session.query(DagModel).filter( DagModel.dag_id.in_([subdag_id, dag_id]), ).filter( DagModel.is_paused.is_(False)).count() self.assertEqual(2, unpaused_dags) DagModel.get_dagmodel(dag.dag_id).set_is_paused(is_paused=True) paused_dags = session.query(DagModel).filter( DagModel.dag_id.in_([subdag_id, dag_id]), ).filter( DagModel.is_paused.is_(True)).count() self.assertEqual(2, paused_dags) session.close()
def test_sync_to_db_default_view(self, mock_now): dag = DAG( 'dag', start_date=DEFAULT_DATE, default_view="graph", ) with dag: DummyOperator(task_id='task', owner='owner1') SubDagOperator(task_id='subtask', owner='owner2', subdag=DAG( 'dag.subtask', start_date=DEFAULT_DATE, )) now = datetime.datetime.utcnow().replace( tzinfo=pendulum.timezone('UTC')) mock_now.return_value = now session = settings.Session() dag.sync_to_db(session=session) orm_dag = session.query(DagModel).filter( DagModel.dag_id == 'dag').one() self.assertIsNotNone(orm_dag.default_view) self.assertEqual(orm_dag.get_default_view(), "graph")
def test_sync_to_db_default_view(self, mock_now): dag = DAG( 'dag', start_date=DEFAULT_DATE, default_view="graph", ) with dag: DummyOperator(task_id='task', owner='owner1') SubDagOperator( task_id='subtask', owner='owner2', subdag=DAG( 'dag.subtask', start_date=DEFAULT_DATE, ) ) now = datetime.datetime.utcnow().replace(tzinfo=pendulum.timezone('UTC')) mock_now.return_value = now session = settings.Session() dag.sync_to_db(session=session) orm_dag = session.query(DagModel).filter(DagModel.dag_id == 'dag').one() self.assertIsNotNone(orm_dag.default_view) self.assertEqual(orm_dag.get_default_view(), "graph")