def test_should_not_trigger_if_current_task_is_last( self, dag, airflow_context, simple_trigger_dag_mock: MagicMock, dag_run): dag_run.conf = {'tasks': ['previous', dag.dag_id]} operator = create_trigger_next_task_dag_operator(dag) operator.execute(airflow_context) simple_trigger_dag_mock.assert_not_called()
def test_should_trigger_next_task(self, dag, airflow_context, simple_trigger_dag_mock: MagicMock, dag_run): dag_run.conf = {'tasks': [dag.dag_id, 'next']} operator = create_trigger_next_task_dag_operator(dag) operator.execute(airflow_context) simple_trigger_dag_mock.assert_called_with('next', dag_run.conf)
def create_dag(dag_id: str = ScienceBeamDagIds.SCIENCEBEAM_CONVERT, default_args: dict = None, schedule_interval=None, macros: ScienceBeamConvertMacros = None, trigger_next: bool = True): if default_args is None: default_args = DEFAULT_ARGS if macros is None: macros = ScienceBeamConvertMacros() dag = DAG(dag_id=dag_id, default_args=default_args, schedule_interval=schedule_interval) convert_results = (create_validate_config_operation( dag=dag, required_props=REQUIRED_PROPS, is_config_valid=macros.is_config_valid) >> create_deploy_sciencebeam_op(dag=dag, macros=macros) >> create_sciencebeam_convert_op(dag=dag, macros=macros)) _ = convert_results >> create_delete_sciencebeam_op(dag=dag) get_output_file_list_results = ( convert_results >> create_get_output_file_list_op(dag=dag, macros=macros)) if trigger_next: _ = get_output_file_list_results >> create_trigger_next_task_dag_operator( dag=dag) return dag
def test_should_not_trigger_without_tasks( self, dag, airflow_context, simple_trigger_dag_mock: MagicMock, dag_run): dag_run.conf = {'other': {}} operator = create_trigger_next_task_dag_operator(dag) operator.execute(airflow_context) simple_trigger_dag_mock.assert_not_called()
def test_should_raise_exception_if_dag_id_not_found( self, dag, airflow_context, simple_trigger_dag_mock: MagicMock, dag_run): dag_run.conf = {'tasks': ['other']} operator = create_trigger_next_task_dag_operator(dag) with pytest.raises(ValueError): operator.execute(airflow_context) simple_trigger_dag_mock.assert_not_called()
def create_dag(): dag = DAG(dag_id=ScienceBeamDagIds.GROBID_BUILD_IMAGE, default_args=DEFAULT_ARGS, schedule_interval=None) _ = (create_grobid_train_validate_config_operation(dag=dag) >> create_build_grobid_image_operator(dag=dag) >> create_trigger_next_task_dag_operator(dag=dag)) return dag
def create_dag(): dag = DAG(dag_id=ScienceBeamDagIds.GROBID_TRAIN_MODEL, default_args=DEFAULT_ARGS, schedule_interval=None) _ = (create_grobid_train_validate_config_operation(dag=dag) >> create_train_grobid_model_operator(dag=dag) >> create_trigger_next_task_dag_operator(dag=dag)) return dag
def create_dag(): dag = DAG(dag_id=ScienceBeamDagIds.SCIENCEBEAM_AUTOCUT_BUILD_IMAGE, default_args=DEFAULT_ARGS, schedule_interval=None) _ = (create_validate_config_operation(dag=dag, required_props=REQUIRED_PROPS) >> create_train_autocut_model_operator(dag=dag) >> create_trigger_next_task_dag_operator(dag=dag)) return dag
def create_dag(): dag = DAG(dag_id=ScienceBeamDagIds.GROBID_TRAIN_EVALUATE_SOURCE_DATASET, default_args=DEFAULT_ARGS, schedule_interval=None) _ = ( create_grobid_train_validate_config_operation(dag=dag) >> create_trigger_sciencebeam_convert_and_evaluate_source_dataset_operator( dag=dag) >> create_trigger_next_task_dag_operator(dag=dag)) return dag
def create_dag(): dag = DAG(dag_id=ScienceBeamDagIds.SCIENCEBEAM_EVALUATION_RESULTS_TO_BQ, default_args=DEFAULT_ARGS, schedule_interval=None) _ = (create_validate_config_operation(dag=dag, required_props=REQUIRED_PROPS) >> create_evaluation_results_to_jsonl_op(dag=dag) >> create_upload_to_bq_op(dag=dag) >> create_remove_jsonl_file_op( dag=dag) >> create_trigger_next_task_dag_operator(dag=dag)) return dag
def create_dag(): dag = DAG( dag_id=ScienceBeamDagIds.SCIENCEBEAM_ADD_MODEL_CONFIG, default_args=DEFAULT_ARGS, schedule_interval=None ) _ = ( create_validate_config_operation(dag=dag, required_props=REQUIRED_PROPS) >> add_model_config_operator(dag=dag) >> create_trigger_next_task_dag_operator(dag=dag) ) return dag
def create_dag(): dag = DAG(dag_id=ScienceBeamDagIds.GROBID_TRAIN_PREPARE, default_args=DEFAULT_ARGS, schedule_interval=None) _ = (create_grobid_train_validate_config_operation(dag=dag) >> create_get_data_operator(dag=dag) >> create_generate_training_data_operator(dag=dag) >> create_auto_annotate_header_operator(dag=dag) >> create_copy_auto_annotated_header_operator(dag=dag) >> create_generate_target_xml_file_list_operator(dag=dag) >> create_generate_header_tei_xml_file_list_operator(dag=dag) >> create_trigger_next_task_dag_operator(dag=dag)) return dag
def create_dag(macros: ScienceBeamEvaluateMacros = None): if macros is None: macros = ScienceBeamEvaluateMacros() dag = DAG(dag_id=ScienceBeamDagIds.SCIENCEBEAM_EVALUATE, default_args=DEFAULT_ARGS, schedule_interval=None) _ = (create_validate_config_operation( dag=dag, required_props=REQUIRED_PROPS, is_config_valid=macros.is_config_valid) >> create_sciencebeam_evaluate_op(dag=dag) >> create_trigger_next_task_dag_operator(dag=dag)) return dag