def test_get_default_config(): td = dagfactory.DagFactory(TEST_DAG_FACTORY) expected = { "default_args": { "owner": "default_owner", "start_date": datetime.date(2018, 3, 1), "end_date": datetime.date(2018, 3, 5), "retries": 1, "retry_delay_sec": 300, }, "concurrency": 1, "max_active_runs": 1, "dagrun_timeout_sec": 600, "default_view": "tree", "orientation": "LR", "schedule_interval": "0 1 * * *", } actual = td.get_default_config() assert actual == expected
def test_doc_md_callable(): td = dagfactory.DagFactory(TEST_DAG_FACTORY) td.generate_dags(globals()) expected_doc_md = globals()['example_dag3'].doc_md assert str(td.get_dag_configs()['example_dag3']['doc_md_python_arguments']) == expected_doc_md
def test_kubernetes_pod_operator_dag(): td = dagfactory.DagFactory(DAG_FACTORY_KUBERNETES_POD_OPERATOR) td.generate_dags(globals()) assert "example_dag" in globals()
def test_generate_dags_invalid(): td = dagfactory.DagFactory(INVALID_DAG_FACTORY) with pytest.raises(Exception): td.generate_dags(globals())
def test_generate_dags_valid(): td = dagfactory.DagFactory(TEST_DAG_FACTORY) td.generate_dags(globals()) assert "example_dag" in globals() assert "example_dag2" in globals() assert "fake_example_dag" not in globals()
def test_generate_dags_invalid(self): td = dagfactory.DagFactory(self.invalid_dag_factory) with pytest.raises(Exception): td.generate_dags(globals())
def test_get_dag_configs_dict(): td = dagfactory.DagFactory(config_filepath=TEST_DAG_FACTORY) assert not set(dagfactory.SYSTEM_PARAMS).issubset(set( td.get_dag_configs()))
def test_generate_dags_valid(self): td = dagfactory.DagFactory(self.test_dag_factory) td.generate_dags(globals()) assert "example_dag" in globals() assert "example_dag2" in globals() assert "fake_example_dag" not in globals()
def test_dagfactory_dict_and_yaml(): error_message = "Either `config_filepath` or `config` should be provided" with pytest.raises(AssertionError, match=error_message): dagfactory.DagFactory(config_filepath=TEST_DAG_FACTORY, config=DAG_FACTORY_CONFIG)
def test_schedule_interval(): td = dagfactory.DagFactory(TEST_DAG_FACTORY) td.generate_dags(globals()) schedule_interval = globals()['example_dag2'].schedule_interval assert schedule_interval is None
def test_get_dag_configs(): td = dagfactory.DagFactory(TEST_DAG_FACTORY) expected = { "example_dag": { "doc_md": "##here is a doc md string", "default_args": { "owner": "custom_owner", "start_date": "2 days" }, "description": "this is an example dag", "schedule_interval": "0 3 * * *", "tasks": { "task_1": { "operator": "airflow.operators.bash_operator.BashOperator", "bash_command": "echo 1", }, "task_2": { "operator": "airflow.operators.bash_operator.BashOperator", "bash_command": "echo 2", "dependencies": ["task_1"], }, "task_3": { "operator": "airflow.operators.bash_operator.BashOperator", "bash_command": "echo 3", "dependencies": ["task_1"], }, }, }, "example_dag2": { "doc_md_file_path": DOC_MD_FIXTURE_FILE, "schedule_interval": "None", "tasks": { "task_1": { "operator": "airflow.operators.bash_operator.BashOperator", "bash_command": "echo 1", }, "task_2": { "operator": "airflow.operators.bash_operator.BashOperator", "bash_command": "echo 2", "dependencies": ["task_1"], }, "task_3": { "operator": "airflow.operators.bash_operator.BashOperator", "bash_command": "echo 3", "dependencies": ["task_1"], }, } }, "example_dag3": { "doc_md_python_callable_name": "mydocmdbuilder", "doc_md_python_callable_file": DOC_MD_PYTHON_CALLABLE_FILE, "doc_md_python_arguments": { "arg1": "arg1", "arg2": "arg2" }, "tasks": { "task_1": { "operator": "airflow.operators.bash_operator.BashOperator", "bash_command": "echo 1", }, } }, } actual = td.get_dag_configs() actual['example_dag2']['doc_md_file_path'] = DOC_MD_FIXTURE_FILE actual['example_dag3'][ 'doc_md_python_callable_file'] = DOC_MD_PYTHON_CALLABLE_FILE assert actual == expected
def test_set_callback_after_loading_config(): td = dagfactory.DagFactory(config=DAG_FACTORY_CONFIG) td.config["default"]["default_args"][ "on_success_callback"] = dagfactory.DagFactory( config=DAG_FACTORY_CONFIG).build_dags td.generate_dags(globals())
def test_generate_dags_with_removal_valid_and_callback(): td = dagfactory.DagFactory(config=DAG_FACTORY_CALLBACK_CONFIG) td.clean_dags(globals()) td.generate_dags(globals())