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())
