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"]
