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