def generate(cls, args):
try:
project_name, project_path = cls.validate_project()
if project_name is None:
cls.parser.error(colored("Invalid project.", 'red'))
if args.type == "dag":
path = "{}/{}/dags".format(project_path, project_name)
if args.m is not None:
path = os.path.join(path, args.m)
if not os.path.exists(path):
cls.parser.error(
colored("The specified module does not exists",
'red'))
Utility.generate_dag_template(project_name, args.n, path)
elif args.type == "module":
path = "{}/{}/dags/{}".format(project_path, project_name,
args.n)
test_path = "{}/tests/{}".format(project_path, args.n)
mod_val = subprocess.call(['mkdir', path])
test_val = subprocess.call(['mkdir', test_path])
if mod_val != 0 or test_val != 0:
cls.parser.error(colored("Unable to generate.", 'red'))
print(colored("Generated successfully.", 'green'))
except Exception as e:
raise AfctlParserException(e)
def test_generate_dag_template(self):
project_name = "tes_project"
path = "/tmp"
dag = "test"
Utility.generate_dag_template(project_name, dag, path)
expected_output = """
from airflow import DAG
from datetime import datetime, timedelta
default_args = {
'owner': 'tes_project',
# 'depends_on_past': ,
# 'start_date': ,
# 'email': ,
# 'email_on_failure': ,
# 'email_on_retry': ,
# 'retries': 0
}
dag = DAG(dag_id='test', default_args=default_args, schedule_interval='@once')
"""
current_output = open(os.path.join('/tmp', 'test_dag.py')).read()
expected_output = expected_output.replace(" ", "")
current_output = current_output.replace(" ", "")
assert expected_output == current_output