Python models.DAG属性代码示例

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_lineage_render(self):
        # tests inlets / outlets are rendered if they are added
        # after initalization
        dag = DAG(
            dag_id='test_lineage_render',
            start_date=DEFAULT_DATE
        )

        with dag:
            op1 = DummyOperator(task_id='task1')

        f1s = "/tmp/does_not_exist_1-{}"
        file1 = File(f1s.format("{{ execution_date }}"))

        op1.inlets.append(file1)
        op1.outlets.append(file1)

        # execution_date is set in the context in order to avoid creating task instances
        ctx1 = {"ti": TI(task=op1, execution_date=DEFAULT_DATE),
                "execution_date": DEFAULT_DATE}

        op1.pre_execute(ctx1)
        self.assertEqual(op1.inlets[0].url, f1s.format(DEFAULT_DATE))
        self.assertEqual(op1.outlets[0].url, f1s.format(DEFAULT_DATE)) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
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()
        self.assertRaises(
            AirflowException,
            _trigger_dag,
            dag_id,
            dag_bag_mock,
            dag_run_mock,
            run_id=None,
            conf=None,
            execution_date=None,
            replace_microseconds=True,
        ) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
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)) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
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 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
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.assertEqual(triggers[0].conf, conf) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_process_bind_param_naive(self):
        """
        Check if naive datetimes are prevented from saving to the db
        """
        dag_id = 'test_process_bind_param_naive'

        # naive
        start_date = datetime.datetime.now()
        dag = DAG(dag_id=dag_id, start_date=start_date)
        dag.clear()

        with self.assertRaises((ValueError, StatementError)):
            dag.create_dagrun(
                run_id=start_date.isoformat,
                state=State.NONE,
                execution_date=start_date,
                start_date=start_date,
                session=self.session
            )
        dag.clear() 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_get_redirect_url(self):
        dag = DAG(DAG_ID, start_date=DEFAULT_DATE)

        with dag:
            task = QuboleOperator(task_id=TASK_ID,
                                  qubole_conn_id=TEST_CONN,
                                  command_type='shellcmd',
                                  parameters="param1 param2",
                                  dag=dag)

        ti = TaskInstance(task=task, execution_date=DEFAULT_DATE)
        ti.xcom_push('qbol_cmd_id', 12345)

        # check for positive case
        url = task.get_extra_links(DEFAULT_DATE, 'Go to QDS')
        self.assertEqual(url, 'http://localhost/v2/analyze?command_id=12345')

        # check for negative case
        url2 = task.get_extra_links(datetime(2017, 1, 2), 'Go to QDS')
        self.assertEqual(url2, '') 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_partition_sensor_error(self, patched_poke):
        patched_poke.return_value = True

        dag = DAG(DAG_ID, start_date=DEFAULT_DATE)

        with self.assertRaises(AirflowException):
            QubolePartitionSensor(
                task_id='test_qubole_partition_sensor',
                poke_interval=1,
                data={
                    "schema": "default",
                    "table": "my_partitioned_table",
                    "columns": [{"column": "month", "values": ["1", "2"]}]
                },
                dag=dag
            ) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_load_templated_yaml(self):
        dag = DAG(dag_id='example_cloudbuild_operator', start_date=TEST_DEFAULT_DATE)
        with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w+t') as build:
            build.writelines("""
            steps:
                - name: 'ubuntu'
                  args: ['echo', 'Hello {{ params.name }}!']
            """)
            build.seek(0)
            body_path = build.name
            operator = CloudBuildCreateBuildOperator(
                body=body_path,
                task_id="task-id", dag=dag,
                params={'name': 'airflow'}
            )
            operator.prepare_template()
            ti = TaskInstance(operator, TEST_DEFAULT_DATE)
            ti.render_templates()
            expected_body = {'steps': [
                {'name': 'ubuntu',
                 'args': ['echo', 'Hello airflow!']
                 }
            ]
            }
            self.assertEqual(expected_body, operator.body) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_templates(self, _):
        dag_id = 'test_dag_id'
        # pylint: disable=attribute-defined-outside-init
        self.dag = DAG(dag_id, default_args={'start_date': DEFAULT_DATE})
        op = CloudDataTransferServiceCreateJobOperator(
            body={"description": "{{ dag.dag_id }}"},
            gcp_conn_id='{{ dag.dag_id }}',
            aws_conn_id='{{ dag.dag_id }}',
            task_id='task-id',
            dag=self.dag,
        )
        ti = TaskInstance(op, DEFAULT_DATE)
        ti.render_templates()
        self.assertEqual(dag_id, getattr(op, 'body')[DESCRIPTION])
        self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
        self.assertEqual(dag_id, getattr(op, 'aws_conn_id')) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_job_delete_with_templates(self, _):
        dag_id = 'test_dag_id'
        args = {'start_date': DEFAULT_DATE}
        self.dag = DAG(dag_id, default_args=args)  # pylint: disable=attribute-defined-outside-init
        op = CloudDataTransferServiceDeleteJobOperator(
            job_name='{{ dag.dag_id }}',
            gcp_conn_id='{{ dag.dag_id }}',
            api_version='{{ dag.dag_id }}',
            task_id=TASK_ID,
            dag=self.dag,
        )
        ti = TaskInstance(op, DEFAULT_DATE)
        ti.render_templates()
        self.assertEqual(dag_id, getattr(op, 'job_name'))
        self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
        self.assertEqual(dag_id, getattr(op, 'api_version')) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_operation_pause_with_templates(self, _):
        dag_id = 'test_dag_id'
        args = {'start_date': DEFAULT_DATE}
        self.dag = DAG(dag_id, default_args=args)  # pylint: disable=attribute-defined-outside-init
        op = CloudDataTransferServicePauseOperationOperator(
            operation_name='{{ dag.dag_id }}',
            gcp_conn_id='{{ dag.dag_id }}',
            api_version='{{ dag.dag_id }}',
            task_id=TASK_ID,
            dag=self.dag,
        )
        ti = TaskInstance(op, DEFAULT_DATE)
        ti.render_templates()
        self.assertEqual(dag_id, getattr(op, 'operation_name'))
        self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
        self.assertEqual(dag_id, getattr(op, 'api_version')) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_operation_resume_with_templates(self, _):
        dag_id = 'test_dag_id'
        args = {'start_date': DEFAULT_DATE}
        self.dag = DAG(dag_id, default_args=args)  # pylint: disable=attribute-defined-outside-init
        op = CloudDataTransferServiceResumeOperationOperator(
            operation_name='{{ dag.dag_id }}',
            gcp_conn_id='{{ dag.dag_id }}',
            api_version='{{ dag.dag_id }}',
            task_id=TASK_ID,
            dag=self.dag,
        )
        ti = TaskInstance(op, DEFAULT_DATE)
        ti.render_templates()
        self.assertEqual(dag_id, getattr(op, 'operation_name'))
        self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
        self.assertEqual(dag_id, getattr(op, 'api_version')) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_instance_start_with_templates(self, _):
        dag_id = 'test_dag_id'
        args = {
            'start_date': DEFAULT_DATE
        }
        self.dag = DAG(dag_id, default_args=args)  # pylint: disable=attribute-defined-outside-init
        op = ComputeEngineStartInstanceOperator(
            project_id='{{ dag.dag_id }}',
            zone='{{ dag.dag_id }}',
            resource_id='{{ dag.dag_id }}',
            gcp_conn_id='{{ dag.dag_id }}',
            api_version='{{ dag.dag_id }}',
            task_id='id',
            dag=self.dag
        )
        ti = TaskInstance(op, DEFAULT_DATE)
        ti.render_templates()
        self.assertEqual(dag_id, getattr(op, 'project_id'))
        self.assertEqual(dag_id, getattr(op, 'zone'))
        self.assertEqual(dag_id, getattr(op, 'resource_id'))
        self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
        self.assertEqual(dag_id, getattr(op, 'api_version')) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DAG [as 别名]
def test_instance_stop_with_templates(self, _):
        dag_id = 'test_dag_id'
        args = {
            'start_date': DEFAULT_DATE
        }
        self.dag = DAG(dag_id, default_args=args)  # pylint: disable=attribute-defined-outside-init
        op = ComputeEngineStopInstanceOperator(
            project_id='{{ dag.dag_id }}',
            zone='{{ dag.dag_id }}',
            resource_id='{{ dag.dag_id }}',
            gcp_conn_id='{{ dag.dag_id }}',
            api_version='{{ dag.dag_id }}',
            task_id='id',
            dag=self.dag
        )
        ti = TaskInstance(op, DEFAULT_DATE)
        ti.render_templates()
        self.assertEqual(dag_id, getattr(op, 'project_id'))
        self.assertEqual(dag_id, getattr(op, 'zone'))
        self.assertEqual(dag_id, getattr(op, 'resource_id'))
        self.assertEqual(dag_id, getattr(op, 'gcp_conn_id'))
        self.assertEqual(dag_id, getattr(op, 'api_version'))