def test_task_dependencies():
input_group = [
{
"operator": "bash",
"task_id": "first_task",
"command": "echo 'This is first task'",
"dependency": "sequential",
"load_from_app_config": "true"
},
{
"operator": "bash",
"task_id": "second_task",
"command": "echo 'This is second task'",
"dependency": "sequential",
"load_from_app_config": "true"
}
]
default_args = {'owner': 'airflow',
'start_date': datetime(2018, 1, 1)
}
dag = DAG("test_dag",
schedule_interval=None,
default_args=default_args)
one_group = LoadDags("").construct_one_group(input_group, "pre_validation", dag)
second_group = LoadDags("").construct_one_group(input_group, "ingestion", dag)
def test_construct_dag():
input_dag = {
"dag_id": "dynamic_trigger_test",
"pre_validation": [
{
"operator": "bash",
"task_id": "first_task",
"command": "echo 'This is first task'",
"dependency": "sequential",
"load_from_app_config": "true"
},
{
"operator": "bash",
"task_id": "second_task",
"command": "echo 'This is second task'",
"dependency": "sequential",
"load_from_app_config": "true"
}
],
"post_validation": [
{
"operator": "bash",
"task_id": "first_task",
"command": "echo 'This is first task'",
"dependency": "sequential",
"load_from_app_config": "true"
},
{
"operator": "bash",
"task_id": "second_task",
"command": "echo 'This is second task'",
"dependency": "sequential",
"load_from_app_config": "true"
}
],
"ingestion": [
{
"operator": "bash",
"task_id": "first_task",
"command": "echo 'This is first task'",
"dependency": "sequential",
"load_from_app_config": "true"
},
{
"operator": "bash",
"task_id": "second_task",
"command": "echo 'This is second task'",
"dependency": "sequential",
"load_from_app_config": "true"
}
]
}
dag_id, dag = LoadDags(input_dag).construct_dag()
task_list = dag.tasks
assert len(task_list) == 6
def test_construct_postgres_task():
input_task = {
"operator": "postgres",
"task_id": "postgres_task",
"conn_id": "postgres_conn_id",
"sql": "select * from table limit 10",
"dependency": "sequential",
"load_from_app_config": "true"
}
default_args = {'owner': 'airflow',
'start_date': datetime(2018, 1, 1)
}
dag = DAG("test_dag",
schedule_interval=None,
default_args=default_args)
operator = LoadDags("").construct_task(input_task, "pre_validation", dag)
assert operator.task_id == "pre_validation_postgres_task"
assert isinstance(operator, PostgresOperator)
assert operator.postgres_conn_id == "postgres_conn_id"
assert operator.sql == "select * from table limit 10"
def test_construct_python_task():
import types
input_task = {
"operator": "python",
"task_id": "python_task",
"function_name": "pow",
"function_kwargs": {"a": 1, "b": 2},
"dependency": "sequential",
"load_from_app_config": "true"
}
default_args = {'owner': 'airflow',
'start_date': datetime(2018, 1, 1)
}
dag = DAG("test_dag",
schedule_interval=None,
default_args=default_args)
operator = LoadDags("").construct_task(input_task, "pre_validation", dag)
assert isinstance(operator, PythonOperator)
assert operator.task_id == "pre_validation_python_task"
assert isinstance(operator.python_callable, types.BuiltinFunctionType)
def test_construct_bash_task():
input_task = {
"operator": "bash",
"task_id": "first_task",
"command": "echo 'This is first task'",
"dependency": "sequential",
"load_from_app_config": "true"
}
default_args = {'owner': 'airflow',
'start_date': datetime(2018, 1, 1)
}
dag = DAG("test_dag",
schedule_interval=None,
default_args=default_args)
operator = LoadDags("").construct_task(input_task, "pre_validation", dag)
assert operator.task_id == "pre_validation_first_task"
assert isinstance(operator, BashOperator)
assert operator.bash_command == "echo 'This is first task'"
import sys
sys.path.insert(0, '/usr/local/airflow')
from airflow.models import DAG
from essentials.load_dag import LoadDags
from essentials.load_json import LoadJSON
input_dag = LoadJSON().load_json_dict()
for single_dag in input_dag:
dag_id, dag = LoadDags(single_dag).construct_dag()
globals()[dag_id] = dag