def oscalize(combined, title): with tempfile.NamedTemporaryFile(mode="w", suffix=".json") as combined_file: combined_file.write(combined) command = "python oscalize.py --title {} {}".format( title, combined_file.name) return "\n".join(ShellTask(command=command, return_all=True).run())
def test_shell_raises_if_no_command_provided(): with Flow(name="test") as f: task = ShellTask()() with pytest.raises(TypeError): with raise_on_exception(): out = f.run()
def test_shell_returns_stderr_as_well(): with Flow(name="test") as f: task = ShellTask()( command="ls surely_a_dir_that_doesnt_exist || exit 0") out = f.run() assert out.is_successful() assert "No such file or directory" in out.result[task].result.decode()
def fit_flow(self): shelltask = ShellTask(return_all=True, log_stderr=True) with Flow("fit") as f: self.fit_shell_task = shelltask(command=self.fit_command) return f
def test_shell_logs_non_zero_exit(caplog): caplog.set_level(level=logging.ERROR, logger="prefect.ShellTask") with Flow(name="test") as f: task = ShellTask()(command="ls surely_a_dir_that_doesnt_exist") out = f.run() assert out.is_failed() assert "Command failed with exit code 2" in caplog.text
def test_shell_inherits_env(): with Flow(name="test") as f: task = ShellTask()(command="echo -n $MYTESTVAR") os.environ["MYTESTVAR"] = "42" out = f.run() assert out.is_successful() assert out.result[task].result == "42"
def test_shell_log_stream_as_info(caplog): with Flow(name="test") as f: ShellTask(stream_output=logging.INFO)(command="echo foo && echo bar") f.run() log_messages = [(r.levelname, r.message) for r in caplog.records] assert ("INFO", "foo") in log_messages and ("INFO", "bar") in log_messages
def test_shell_task_accepts_env(): with Flow(name="test") as f: task = ShellTask()(command="echo -n $MYTESTVAR", env=dict(MYTESTVAR="test")) out = f.run() assert out.is_successful() assert out.result[task].result == "test"
def build(self): tasks = [] with open('sobjects.txt') as sobjects_file: for sobject in sobjects_file.readlines(): target = sobject.strip() env = { 'SALESFORCE_USER': os.environ.get('SALESFORCE_USER'), 'SALESFORCE_PASSWORD': os.environ.get('SALESFORCE_PASSWORD'), 'SALESFORCE_TOKEN': os.environ.get('SALESFORCE_TOKEN'), 'POSTGRES_USER': os.environ.get('POSTGRES_USER'), 'POSTGRES_PASSWORD': os.environ.get('POSTGRES_PASSWORD'), 'SALESFORCE_SOBJECT': target, 'POSTGRES_TABLE': 'sf_' + snakecase(target), } command = 'java -jar /bin/embulk run /work/salesforce.yml.liquid -c /work/diff/{}.diff.yml'.format( target) task = ShellTask(name=target, command=command, log_stdout=True, return_all=True, env=env) tasks.append(task) with Flow("build source", tasks=tasks) as f: pass out = f.run()
def fit(self, train_filename, model_filename=None, verbose=False): self.train_filename = train_filename if model_filename is not None: self.model_filename = model_filename self.model_command = """-{} {}""".format( self.prefixes.get("dump_model"), self.model_filename) self.train_command = """-{} {}""".format(self.prefixes.get("train"), self.train_filename) self.command = """{} {} {} {}""".format( self.prefixes.get("weka"), self.algorithm_command, self.train_command, self.model_command, ) shell = ShellTask() with Flow("fit") as f: fit = shell(command=self.command) status = f.run() if verbose and status.is_successful(): print(status.result[fit].result.decode("utf-8")) return status
def test_shell_task_env_can_be_set_at_init(): with Flow(name="test") as f: task = ShellTask(env=dict(MYTESTVAR="test"))( command="echo -n $MYTESTVAR") out = f.run() assert out.is_successful() assert out.result[task].result == "test"
def test_shell_task_accepts_helper_script(): helper = "cd ~" task = ShellTask() with Flow("test") as f: res = task(command="ls", helper_script=helper) out = f.run() assert out.is_successful()
def test_shell_adds_newline_to_helper_script(): helper = "cd ~" task = ShellTask(helper_script=helper) with Flow("test") as f: res = task(command="ls") out = f.run() assert out.is_successful()
def test_shell_logs_error_on_non_zero_exit(caplog): with Flow(name="test") as f: task = ShellTask()(command="ls surely_a_dir_that_doesnt_exist") out = f.run() assert out.is_failed() assert "ERROR prefect.Task: ShellTask:shell.py" in caplog.text assert "surely_a_dir_that_doesnt_exist" in caplog.text assert "No such file or directory" in caplog.text
def markdown(combined, dest_dir): with tempfile.NamedTemporaryFile(mode="w", suffix=".json") as combined_file: combined_file.write(combined) command = "python component_report.py {} {}".format( combined_file.name, Path(dest_dir) / "markdown") return "\n".join(ShellTask(command=command, return_all=True).run())
def test_shell_respects_stream_output(caplog, stream_output): with Flow(name="test") as f: ShellTask(stream_output=stream_output)( command="echo foo && echo bar", ) f.run() stdout_in_log = "foo" in caplog.messages and "bar" in caplog.messages assert stdout_in_log == stream_output
def test_shell_logs_non_zero_exit(caplog): with Flow(name="test") as f: task = ShellTask()(command="ls surely_a_dir_that_doesnt_exist") out = f.run() assert out.is_failed() error_log = [c for c in caplog.records if c.levelname == "ERROR"] assert len(error_log) == 1 assert error_log[0].name == "prefect.ShellTask" assert "Command failed" in error_log[0].message
def test_shell_log_stream_as_info_string_input(caplog, stream_output): with Flow(name="test") as f: ShellTask(stream_output=stream_output)(command="echo foo && echo bar") f.run() log_messages = [(r.levelname, r.message) for r in caplog.records] assert (stream_output, "foo") in log_messages and ( stream_output, "bar", ) in log_messages
def test_shell_attaches_result_to_failure(caplog): def assert_fail_result(task, state): assert ( state.result == "foo" "ls: surely_a_dir_that_doesnt_exist: No such file or directory") with Flow(name="test") as f: task = ShellTask(on_failure=assert_fail_result)( command="echo foo && ls surely_a_dir_that_doesnt_exist") out = f.run() assert out.is_failed()
def combine(components): paths = [] for component in components: component_file = tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") component_file.write(component) paths.append(component_file.name) component_file.close() paths_arg = " ".join(paths) command = "python ssp.py combine {}".format(paths_arg) return "\n".join(ShellTask(command=command, return_all=True).run())
def test_shell_logs_stderr_on_non_zero_exit(caplog): caplog.set_level(level=logging.ERROR, logger="prefect.ShellTask") with Flow(name="test") as f: task = ShellTask( log_stderr=True, return_all=True)(command="ls surely_a_dir_that_doesnt_exist") out = f.run() assert out.is_failed() assert len(caplog.records) == 2 assert "Command failed" in caplog.records[0].message assert "No such file or directory" in caplog.records[1].message assert "surely_a_dir_that_doesnt_exist" in caplog.records[1].message
def test_shell_sources_helper_script_correctly(): helper = """ say_hi() { echo $1 } """ task = ShellTask(helper_script=helper) with Flow(name="test") as f: res = task(command="say_hi chris") out = f.run() assert out.is_successful() assert out.result[res].result == "chris"
def test_shell_logs_stderr_on_non_zero_exit(caplog): with Flow(name="test") as f: task = ShellTask( log_stderr=True, return_all=True)(command="ls surely_a_dir_that_doesnt_exist") out = f.run() assert out.is_failed() error_log = [c for c in caplog.records if c.levelname == "ERROR"] assert len(error_log) == 2 assert error_log[0].name == "prefect.ShellTask" assert "Command failed" in error_log[0].message assert "No such file or directory" in error_log[1].message assert "surely_a_dir_that_doesnt_exist" in error_log[1].message
def test_shell_initializes_and_runs_multiline_cmd(): cmd = """ for line in $(printenv) do echo "${line#*=}" done """ with Flow(name="test") as f: task = ShellTask()(command=cmd, env={key: "test" for key in "abcdefgh"}) out = f.run() assert out.is_successful() lines = out.result[task].result.decode().split("\n") test_lines = [l for l in lines if l == "test"] assert len(test_lines) == 8
def predict(self, test_filename, prediction_filename=None, verbose=False, **kwargs): self.test_filename = test_filename if prediction_filename is not None: self.prediction_filename = prediction_filename else: self.prediction_filename = 0 # Weka default self.model_command = """-{} {}""".format( self.prefixes.get("load_model"), self.model_filename) self.test_command = """-{} {}""".format(self.prefixes.get("test"), self.test_filename) self.predict_command = ( """-classifications weka.classifiers.evaluation.output.prediction.CSV""" ) self.command = """{} {} {} {} {}""".format( self.prefixes.get("weka"), self.prefixes.get("J48"), self.test_command, self.model_command, self.predict_command, ) shell = ShellTask() with Flow("predict") as f: predict = shell(command=self.command) status = f.run() if verbose and status.is_successful(): output = status.result[predict].result.decode("utf-8") df = self._parse_output(output) return df else: return status
def test_shell_initializes_and_runs_multiline_cmd(): cmd = """ TEST=$(cat <<-END This is line one This is line two This is line three boom! END ) for i in $TEST do echo $i done""" with Flow(name="test") as f: task = ShellTask()(command=cmd, env={key: "test" for key in "abcdefgh"}) out = f.run() assert out.is_successful() assert out.result[task].result == "boom!"
def test_shell_task_handles_multiline_commands(): with tempfile.TemporaryDirectory() as tempdir: cmd = """ cd {} for file in $(ls) do cat $file done """.format(tempdir) with open(tempdir + "/testfile.txt", "w") as f: f.write("this is a test") with Flow(name="test") as f: task = ShellTask()(command=cmd) out = f.run() assert out.is_successful() assert out.result[task].result == "this is a test"
def test_basic_trigger_dag_triggers(self, airflow_settings): task = AirflowTriggerDAG(dag_id="tutorial", execution_date="1986-09-20", env=airflow_settings) check_task = ShellTask( command="airflow list_dag_runs tutorial", helper_script=task.helper_script, env=airflow_settings, ) with Flow(name="tutorial") as flow: res = check_task(upstream_tasks=[task]) flow_state = flow.run() assert flow_state.is_successful() check_state = flow_state.result[res] assert check_state.is_successful() # check CLI output assert "manual__1986-09-20T00:00:00+00:00" in check_state.result assert "running" in check_state.result assert "1986-09-20T00:00:00+00:00" in check_state.result
def test_shell_runs_other_shells(): with Flow(name="test") as f: task = ShellTask(shell="zsh")(command="echo -n $ZSH_NAME") out = f.run() assert out.is_successful() assert out.result[task].result == "zsh"
def test_shell_initializes_and_multiline_output_optionally_returns_all_lines(): with Flow(name="test") as f: task = ShellTask(return_all=True)(command="echo -n 'hello world\n42'") out = f.run() assert out.is_successful() assert out.result[task].result == ["hello world", "42"]