def setUp(self): if ".airflow" not in os.listdir("."): os.mkdir(".airflow") if "dags" not in os.listdir(".airflow"): os.mkdir(".airflow/dags") if "temp" not in os.listdir(".airflow"): os.mkdir(".airflow/temp") self.write_dir = os.getcwd() + "/.airflow/temp/" with open(self.write_dir + "foo.sh", "w") as f: f.write("printf 'a,b,c\n1,2,3' >> {0}/foo.csv".format( self.write_dir)) with open(".airflow/temp/foo2.sh", "w") as f: f.write("printf 'Success!' >> {0}/foo2.csv".format(self.write_dir)) self.dep_sh = Depositor("TestDepositor", ".airflow/temp/foo.sh", ".airflow/temp/foo.csv") self.trans_sh = Transform("TestTransform", ".airflow/temp/foo2.sh", ".airflow/temp/foo.csv", ".airflow/temp/foo2.csv", requirements=[self.dep_sh]) self.dag = DAG('dag', default_args=default_args) orchestration.configure(init=True) orchestration.write_airflow_string([self.dep_sh, self.trans_sh], "./.airflow/dags/airscooter_dag.py")
def setUp(self): self.dep_sh = Depositor("TestDepositor", "foo.sh", "foo.csv") self.trans_sh = Transform("TestTransform", "foo2.sh", "foo.csv", "foo2.csv", requirements=[self.dep_sh]) self.dag = DAG('dag', default_args=default_args)
def setUp(self): if ".airflow" not in os.listdir("."): os.mkdir(".airflow") if "dags" not in os.listdir(".airflow"): os.mkdir(".airflow/dags") if "temp" not in os.listdir(".airflow"): os.mkdir(".airflow/temp") self.write_dir = os.getcwd() + "/.airflow/temp/" with open(self.write_dir + "foo.py", "w") as f: # noinspection SqlNoDataSourceInspection,SqlDialectInspection f.write( """import pandas\nwith open('{0}/foo.csv', 'w') as f: f.write('1,2,3')""" .format(self.write_dir)) with open(".airflow/temp/foo2.py", "w") as f: # noinspection SqlNoDataSourceInspection,SqlDialectInspection f.write( """with open('{0}/foo2.csv', 'w') as f: f.write('1,2,3')""". format(self.write_dir)) import subprocess subprocess.run([ "cp", "./fixtures/test_nb.ipynb", os.getcwd() + "/.airflow/temp/foo3.ipynb" ]) self.py_dep_sh = Depositor("TestDepositor", ".airflow/temp/foo.py", ".airflow/temp/foo.csv") self.py_trans_sh = Transform("TestPyTransform", ".airflow/temp/foo2.py", ".airflow/temp/foo.csv", ".airflow/temp/foo2.csv", requirements=[self.py_dep_sh]) self.ipynb_trans_sh = Transform("TestJupyterTransform", ".airflow/temp/foo3.ipynb", ".airflow/temp/foo.csv", ".airflow/temp/foo3.csv", requirements=[self.py_dep_sh]) self.dummy_trans_sh = Transform("TestJupyterTransform", ".airflow/temp/foo3.ipynb", ".airflow/temp/nan.csv", ".airflow/temp/nan2.csv", requirements=[self.py_dep_sh], dummy=True) self.dag = DAG('dag', default_args=default_args) orchestration.configure(init=True) orchestration.write_airflow_string( [self.py_dep_sh, self.py_trans_sh, self.ipynb_trans_sh], "./.airflow/dags/airscooter_dag.py")
def test_transform_datafy(): dep_sh = Depositor("TestDepositor", "foo.sh", "foo.csv") trans_sh = Transform("TestTransform", "foo.sh", "foo.csv", "foo2.csv", requirements=[dep_sh]) t = trans_sh.datafy() assert all([ key in t for key in ['name', 'filename', 'input', 'output', 'type', 'requirements'] ]) assert t['requirements'] == ["TestDepositor"]
def setUp(self): self.dep_sh = Depositor("TestDepositor", "foo.sh", "foo.csv") self.trans_sh = Transform("TestTransform", "foo2.sh", "foo.csv", "foo2.csv", requirements=[self.dep_sh]) self.dag = DAG('dag', default_args=default_args) if ".airflow" not in os.listdir("."): os.mkdir(".airflow") if "dags" not in os.listdir(".airflow"): os.mkdir(".airflow/dags") set_up_files()
class TestShellDepositor(unittest.TestCase): def setUp(self): set_up_files() self.dep_sh = Depositor("TestTransform", "foo.sh", "foo.csv") self.dag = DAG('dag', default_args=default_args) def test_this(self): # noinspection PyUnresolvedReferences from airflow.operators.bash_operator import BashOperator d = self.dep_sh.as_airflow_string() # dag -> self.dag d = replace_once_right(d, "dag", "self.dag") bash_operator = eval(d) assert bash_operator.bash_command.rstrip() == "echo HELLO" def tearDown(self): tear_files_down()
def setUp(self): set_up_files() self.dep_sh = Depositor("TestTransform", "foo.sh", "foo.csv") self.dag = DAG('dag', default_args=default_args)