def test_extra_serialized_field_and_multiple_operator_links(self): """ Assert extra field exists & OperatorLinks defined in Plugins and inbuilt Operator Links. This tests also depends on GoogleLink() registered as a plugin in tests/plugins/test_plugin.py The function tests that if extra operator links are registered in plugin in ``operator_extra_links`` and the same is also defined in the Operator in ``BaseOperator.operator_extra_links``, it has the correct extra link. """ test_date = datetime(2019, 8, 1) dag = DAG(dag_id='simple_dag', start_date=test_date) CustomOperator(task_id='simple_task', dag=dag, bash_command=["echo", "true"]) serialized_dag = SerializedDAG.to_dict(dag) assert "bash_command" in serialized_dag["dag"]["tasks"][0] dag = SerializedDAG.from_dict(serialized_dag) simple_task = dag.task_dict["simple_task"] assert getattr(simple_task, "bash_command") == ["echo", "true"] ######################################################### # Verify Operator Links work with Serialized Operator ######################################################### # Check Serialized version of operator link only contains the inbuilt Op Link assert serialized_dag["dag"]["tasks"][0]["_operator_extra_links"] == [ {'tests.test_utils.mock_operators.CustomBaseIndexOpLink': {'index': 0}}, {'tests.test_utils.mock_operators.CustomBaseIndexOpLink': {'index': 1}}, ] # Test all the extra_links are set assert set(simple_task.extra_links) == { 'BigQuery Console #1', 'BigQuery Console #2', 'airflow', 'github', 'google', } ti = TaskInstance(task=simple_task, execution_date=test_date) ti.xcom_push('search_query', ["dummy_value_1", "dummy_value_2"]) # Test Deserialized inbuilt link #1 custom_inbuilt_link = simple_task.get_extra_links(test_date, "BigQuery Console #1") assert 'https://console.cloud.google.com/bigquery?j=dummy_value_1' == custom_inbuilt_link # Test Deserialized inbuilt link #2 custom_inbuilt_link = simple_task.get_extra_links(test_date, "BigQuery Console #2") assert 'https://console.cloud.google.com/bigquery?j=dummy_value_2' == custom_inbuilt_link # Test Deserialized link registered via Airflow Plugin google_link_from_plugin = simple_task.get_extra_links(test_date, GoogleLink.name) assert "https://www.google.com" == google_link_from_plugin
def test_extra_serialized_field_and_operator_links(self): """ Assert extra field exists & OperatorLinks defined in Plugins and inbuilt Operator Links. This tests also depends on GoogleLink() registered as a plugin in tests/plugins/test_plugin.py The function tests that if extra operator links are registered in plugin in ``operator_extra_links`` and the same is also defined in the Operator in ``BaseOperator.operator_extra_links``, it has the correct extra link. """ test_date = datetime(2019, 8, 1) dag = DAG(dag_id='simple_dag', start_date=test_date) CustomOperator(task_id='simple_task', dag=dag, bash_command="true") serialized_dag = SerializedDAG.to_dict(dag) self.assertIn("bash_command", serialized_dag["dag"]["tasks"][0]) dag = SerializedDAG.from_dict(serialized_dag) simple_task = dag.task_dict["simple_task"] self.assertEqual(getattr(simple_task, "bash_command"), "true") ######################################################### # Verify Operator Links work with Serialized Operator ######################################################### # Check Serialized version of operator link only contains the inbuilt Op Link self.assertEqual( serialized_dag["dag"]["tasks"][0]["_operator_extra_links"], [{ 'tests.test_utils.mock_operators.CustomOpLink': {} }], ) # Test all the extra_links are set self.assertCountEqual(simple_task.extra_links, ['Google Custom', 'airflow', 'github', 'google']) ti = TaskInstance(task=simple_task, execution_date=test_date) ti.xcom_push('search_query', "dummy_value_1") # Test Deserialized inbuilt link custom_inbuilt_link = simple_task.get_extra_links( test_date, CustomOpLink.name) self.assertEqual( 'http://google.com/custom_base_link?search=dummy_value_1', custom_inbuilt_link) # Test Deserialized link registered via Airflow Plugin google_link_from_plugin = simple_task.get_extra_links( test_date, GoogleLink.name) self.assertEqual("https://www.google.com", google_link_from_plugin)
def make_simple_dag(): """Make very simple DAG to verify serialization result.""" dag = DAG( dag_id='simple_dag', default_args={ "retries": 1, "retry_delay": timedelta(minutes=5), "depends_on_past": False, }, start_date=datetime(2019, 8, 1), ) BaseOperator(task_id='simple_task', dag=dag, owner='airflow') CustomOperator(task_id='custom_task', dag=dag) return {'simple_dag': dag}
def make_simple_dag(): """Make very simple DAG to verify serialization result.""" dag = DAG(dag_id='simple_dag', default_args={ "retries": 1, "retry_delay": timedelta(minutes=5), "depends_on_past": False, }, start_date=datetime(2019, 8, 1), is_paused_upon_creation=False, access_control={"test": {"can_dag_read", "can_dag_edit"}}) BaseOperator(task_id='simple_task', dag=dag, owner='airflow') CustomOperator(task_id='custom_task', dag=dag) return {'simple_dag': dag}
def make_simple_dag(): """Make very simple DAG to verify serialization result.""" with DAG( dag_id='simple_dag', default_args={ "retries": 1, "retry_delay": timedelta(minutes=5), "depends_on_past": False, }, start_date=datetime(2019, 8, 1), is_paused_upon_creation=False, ) as dag: CustomOperator(task_id='custom_task') BashOperator(task_id='bash_task', bash_command='echo {{ task.task_id }}', owner='airflow') return {'simple_dag': dag}
def make_simple_dag(): """Make very simple DAG to verify serialization result.""" with DAG(dag_id='simple_dag', default_args={ "retries": 1, "retry_delay": timedelta(minutes=5), "depends_on_past": False, }, start_date=datetime(2019, 8, 1), is_paused_upon_creation=False, access_control={"test_role": {"can_read", "can_edit"}}) as dag: CustomOperator(task_id='custom_task') BashOperator(task_id='bash_task', bash_command='echo {{ task.task_id }}', owner='airflow', executor_config={"pod_override": executor_config_pod}) return {'simple_dag': dag}