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_can_push_with_tags(self):
self.tag_repo(['1.0.0', '1', '1.0', '1.0.0+20200228.blue-ivory'])
checkout = SemanticCheckoutTask(upstream_repos=dict(
abc=f'sgr://{remote_name}/abc/1234?tag=1', ), )
workspaces = checkout.run()
push = PushRepoTask(workspaces=workspaces, )
self.repo.commit()
runner = TaskRunner(task=push)
tags_edge = Edge(Task(), push, key='sgr_tags')
tag_state = Success(result=ConstantResult(value=dict(
abc=['foo', 'bar', 'tag1_w_upstream'])))
with raise_on_exception():
with prefect.context():
state = runner.run(upstream_states={tags_edge: tag_state})
if state.is_failed():
print(state)
self.fail()
self.assertCountEqual(
self.repo.head.get_tags(),
['HEAD', 'foo', 'bar', 'tag1_w_upstream'])
def test_retries_cache_parameters_as_well(self, executor):
with Flow(name="test") as f:
a = Parameter("a")
b = ReturnTask(max_retries=1, retry_delay=datetime.timedelta(0))
a_res = a()
b_res = b(a_res)
first_state = FlowRunner(flow=f).run(executor=executor,
parameters=dict(a=1),
return_tasks=f.tasks)
assert first_state.is_running()
a_state = first_state.result[a_res]
a_state.result = (
NoResult
) # remove the result to see if the cached results are picked up
b_state = first_state.result[b_res]
b_state.cached_inputs = dict(x=Result(2)) # artificially alter state
with raise_on_exception(): # without caching we'd expect a KeyError
second_state = FlowRunner(flow=f).run(
executor=executor,
return_tasks=[b_res],
task_states=first_state.result)
assert isinstance(second_state, Success)
assert second_state.result[b_res].result == 1
def test_can_clone_with_hourly_prerelease_tags(self):
self.tag_repo([
'1.0.0', '1', '1.0', '1.0.0+20200228.blue-ivory', '1.0.1',
'1.0.1+20200307.pink-bear', '1-hourly', '1.0-hourly',
'1.0.2-hourly.1+20200301.blue-ivory',
'1.0.2-hourly.2+20200301.green-monster'
])
checkout = SemanticCheckoutTask(upstream_repos=dict(
abc='sgr:///abc/1234?tag=1.1-hourly', ), )
runner = TaskRunner(task=checkout)
with raise_on_exception():
with prefect.context():
state = runner.run()
if state.is_failed():
print(state)
self.fail()
workspaces = state.result
workspace = workspaces['abc']
self.assertEqual(
workspace['version'],
Version('1.0.2-hourly.2+20200301.green-monster'))
def test_raise_on_exception_raises_basic_prefect_signal():
flow = Flow(name="test")
flow.add_task(BusinessTask())
with pytest.raises(prefect.engine.signals.FAIL) as error:
with raise_on_exception():
flow.run()
assert "needs more blockchain!" in str(error.value)
def run(
self,
small_area_period_built_filepath: Path,
ber_archetypes_filepath: Path,
small_area_geometries_filepath: Path,
output_filepath: Path,
) -> gpd.GeoDataFrame:
"""Run EstimateSmallAreaDemand flow.
Args:
small_area_period_built_filepath (Path): Path to Clean Small Area Period
Built Data
ber_archetypes_filepath (Path): Path to BER Archetypal Energy Averages
small_area_geometries_filepath (Path): Path to Clean Small Area Geometries
output_filepath (Path): Path to a Small Area Demand Estimate
"""
with raise_on_exception():
state = flow.run(parameters=dict(
sa_periods_built_fpath=small_area_period_built_filepath,
ber_archetype_averages_fpath=ber_archetypes_filepath,
sa_geometries_fpath=small_area_geometries_filepath,
), )
result = state.result[sa_demand_with_geometries].result
result.to_parquet(output_filepath)
def test_flow_runner_uses_user_provided_executor():
t = SuccessTask()
with Flow(name="test") as f:
result = t()
with raise_on_exception():
with pytest.raises(NotImplementedError):
FlowRunner(flow=f).run(executor=Executor())
def test_can_semantic_bump_prerelease(self):
semantic_bump = SemanticBumpTask()
runner = TaskRunner(task=semantic_bump)
edge = Edge(Task(), semantic_bump, key='workspaces')
upstream_state = Success(result=ConstantResult(value=dict(abc=dict(
repo_uri='sgr:///abc/1234?tag=1-hourly',
version=Version('1.0.1-hourly.4+2021-03-08.zip-fur'),
))))
date = datetime.utcnow()
flow_run_name = 'testflow1'
with raise_on_exception():
with prefect.context(date=date, flow_run_name=flow_run_name):
state = runner.run(upstream_states={edge: upstream_state})
if state.is_failed():
print(state)
self.fail()
self.assertEqual(
state.result,
dict(abc=[
'1-hourly', '1.0-hourly',
f'1.0.1-hourly.5+{date:%Y-%m-%dT%H}.{date:%M}.{flow_run_name}'
]))
def run(
self,
input_filepath: Path,
sa_glossary_filepath: Path,
postcodes_filepath: Path,
sa_geometries_filepath: Path,
output_filepath_period_built: Path,
output_filepath_boilers: Path,
) -> None:
"""Run local flow.
Args:
input_filepath (Path): Path to Small Area Statistics Raw Data
sa_glossary_filepath (Path): Path to Small Area Statistics Glossary
postcodes_filepath (Path): Path to Postcode Geometries Data
sa_geometries_filepath (Path): Path to Small Area Geometries Data
output_filepath_period_built (Path): Path to Small Area Period Built Stats
output_filepath_boilers (Path): Path to Small Area Period Boiler Stats
"""
with raise_on_exception():
state = self.flow.run(parameters=dict(
sa_stats_fpath=input_filepath,
sa_glossary_fpath=sa_glossary_filepath,
postcode_geometries_fpath=postcodes_filepath,
sa_geometries_fpath=sa_geometries_filepath,
), )
period_built = state.result[clean_year_built].result
boilers = state.result[clean_boiler_stats].result
period_built.to_parquet(output_filepath_period_built)
boilers.to_parquet(output_filepath_boilers)
def etl_flow_state(monkeypatch: MonkeyPatch, mock_data_dir: Mock) -> State:
"""Run etl flow with dummy test data.
Args:
monkeypatch (MonkeyPatch): Pytest fixture to mock out objects s.a. PrefectSecret
mock_data_dir (Mock): See https://docs.pytest.org/en/stable/tmpdir.html
Returns:
[State]: A Prefect State object containing flow run information
"""
from drem.etl import residential
# Mock out PrefectSecret as no secrets are required in CI (data already downloaded)
monkeypatch.setattr(
residential.PrefectSecret,
"run",
mock_prefectsecret_run,
)
# Mock out task load as CI doesn't need flow outputs on-disk
monkeypatch.setattr(
residential.Download,
"run",
mock_task_run,
)
monkeypatch.setattr(
residential.download.BERPublicsearch,
"run",
mock_task_run,
)
with raise_on_exception():
state = residential.flow.run()
return state
def test_not_forced_with_prefect_cache(mocker, temp_url):
# Test that the prefect cache is used when force=False
run_method = mocker.patch(
'iguazu.Task.run',
side_effect=['result1', Exception('should not be called')])
task = Task(force=False)
with Flow('test_not_forced_with_prefect_cache') as flow:
task()
with prefect.context(caches={}):
flow_state1 = flow.run(run_on_schedule=True)
cache = prefect.context.caches
result1 = list(flow_state1.result.values())[0].result
with raise_on_exception(), prefect.context(
caches=cache): # Note: cache is from previous run
flow_state2 = flow.run(run_on_schedule=True)
result2 = list(flow_state2.result.values())[0].result
assert result1 == 'result1'
assert result2 == 'result1'
assert run_method.call_count == 1
def test_map_composition(executor):
ll = ListTask()
a = AddTask()
with Flow(name="test") as f:
r1 = a.map(ll)
r2 = a.map(r1)
with raise_on_exception():
s = f.run(executor=executor)
m1 = s.result[r1]
m2 = s.result[r2]
assert s.is_successful()
assert m1.is_mapped()
assert m2.is_mapped()
assert isinstance(m1.map_states, list)
assert all(s.is_successful() for s in m1.map_states)
assert len(m1.map_states) == 3
assert m1.result == [2, 3, 4]
assert isinstance(m2.map_states, list)
assert all(s.is_successful() for s in m2.map_states)
assert len(m2.map_states) == 3
assert m2.result == [3, 4, 5]
def test_function_order(mocker):
manager = mocker.Mock()
funcs = {
'contexts': tuple(),
'prepare_inputs': dict(),
'handle_outputs': None,
}
for name, rv in funcs.items():
mocked_func = mocker.patch(f'iguazu.core.tasks.ManagedTask.{name}',
return_value=rv)
manager.attach_mock(mocked_func, name)
with Flow('test_function_order') as flow:
instance = ManagedTask()
instance()
with raise_on_exception(), prefect.context(caches={}):
flow.run()
expected_calls = [
mocker.call.contexts(),
mocker.call.prepare_inputs(),
mocker.call.handle_outputs(None),
]
manager.assert_has_calls(expected_calls, any_order=False)
def test_graceful_exceptions_fail(temp_url):
task = TaskWithException(graceful_exceptions=())
with Flow('test_graceful_exceptions_fail') as flow:
task(value=-1)
with raise_on_exception(), prefect.context(caches={}), pytest.raises(CustomException):
flow.run()
def test_flow_runner_uses_default_executor_on_flow_if_present():
t = SuccessTask()
with Flow(name="test", executor=Executor()) as flow:
result = t()
with raise_on_exception():
with pytest.raises(NotImplementedError):
FlowRunner(flow=flow).run()
def execute(flow: Flow) -> state:
"""
Returns:
state: (state) state of league flow
"""
with raise_on_exception():
executor = DaskExecutor(address=os.getenv("WORKER_ADDRESS"))
state = flow.run()
return state
def test_auto_manage_dataframe_incorrect_type(local_file, caplog):
task = TaskWithInputDataFrame()
with Flow('test_auto_manage_dataframe_default_string') as flow:
file = prefect.Parameter('local_file', default=123456)
task(input_one=file)
with pytest.raises(ValueError), caplog.at_level('FATAL', logger='prefect'):
with raise_on_exception(), prefect.context(caches={}):
flow.run()
def test_raise_on_exception_plays_well_with_context():
flow = Flow(name="test")
flow.add_task(MathTask())
try:
assert "raise_on_exception" not in prefect.context
with raise_on_exception():
assert "raise_on_exception" in prefect.context
flow.run()
except ZeroDivisionError:
assert "raise_on_exception" not in prefect.context
pass
def test_graceful_exceptions_handle_failed(mocker, temp_url):
graceful_fail_task_method = mocker.patch('iguazu.core.tasks.Task._graceful_fail',
side_effect=[ENDRUN(state=Finished())])
task = TaskWithException(graceful_exceptions=(CustomException, ))
with Flow('test_graceful_exceptions_handle_failed') as flow:
task(value=+1)
with raise_on_exception(), prefect.context(caches={}), pytest.raises(ValueError):
flow.run()
graceful_fail_task_method.assert_not_called()
def test_pandas_context_raise(temp_url):
initial_mode = pd.get_option('mode.chained_assignment')
with Flow('test_pandas_context_error') as flow:
task = PandasModeSpy(pandas_chained_assignment='raise')
task()
with pytest.raises(pandas.core.common.SettingWithCopyError):
with raise_on_exception(), prefect.context(caches={}):
flow.run()
assert pd.get_option('mode.chained_assignment') == initial_mode
def test_switch_works_with_raise_on_exception():
@prefect.task
def return_b():
return "b"
tasks = {let: prefect.Task(name=let) for let in "abcde"}
with Flow(name="test") as flow:
res = switch(return_b, tasks)
with raise_on_exception():
flow_state = flow.run()
def test_bad_graceful_implementation(mocker, temp_url):
graceful_fail_task_method = mocker.patch('iguazu.core.tasks.Task._graceful_fail',
return_value=['something'])
task = TaskWithException(graceful_exceptions=(CustomException, ))
with Flow('test_bad_graceful_implementation') as flow:
task(value=-1)
with raise_on_exception(), prefect.context(caches={}), pytest.raises(RuntimeError):
flow.run()
graceful_fail_task_method.assert_called_once()
def test_mapping_over_constants():
@prefect.task
def add_one(x):
return x + 1
with Flow("constants") as f:
output = add_one.map(x=[1, 2, 3, 4])
with raise_on_exception():
flow_state = f.run()
assert flow_state.is_successful()
assert flow_state.result[output].result == [2, 3, 4, 5]
def run(self, input_filepath: Path, output_filepath: Path) -> None:
"""Run flow.
Args:
input_filepath (Path): Path to input data
output_filepath (Path): Path to output data
"""
with raise_on_exception():
state = self.flow.run(ber_fpath=input_filepath)
result = state.result[bin_year_built_into_census_categories].result
result.to_parquet(output_filepath)
def run(self, input_filepath: Path, output_filepath: Path) -> pd.DataFrame:
"""Run Flow.
Args:
input_filepath (Path): Path to Clean BER Data
output_filepath (Path): Path to BER archetypes
"""
with raise_on_exception():
state = self.flow.run(parameters=dict(ber_fpath=input_filepath))
result = state.result[ber_archetypes].result
result.to_parquet(output_filepath)
def test_raise_on_exception_ignores_all_prefect_signals(signal):
flow = Flow(name="test")
@prefect.task
def raise_signal():
if (prefect.context.get("task_loop_count", 1) < 2
and prefect.context.get("task_run_count", 1) < 2
and signal.__name__ != "PAUSE"):
raise signal("my message")
flow.add_task(raise_signal)
with raise_on_exception():
flow_state = flow.run()
def test_auto_manage_dataframe_init_overwrite(local_file):
task = TaskWithInputDataFrame(input_one_key='/bar')
with Flow('test_auto_manage_dataframe_init_overwrite') as flow:
file = prefect.Parameter('local_file', default=local_file)
task(input_one=file)
with raise_on_exception(), prefect.context(caches={}):
flow_state = flow.run()
result = list(flow_state.result.values())[0].result
df_bar = pd.read_hdf(local_file.file, '/bar')
assert isinstance(df_bar, pd.DataFrame)
tm.assert_equal(result, df_bar.mean())
def test_flatmap_reduced_result(self, executor):
# flatmap over a reduced flattened mapped task
ll = ListTask()
nest = NestTask()
a = AddTask()
with Flow(name="test") as f:
nested = nest.map(ll())
nested2 = nest(flatten(nested))
x = a.map(flatten(nested2))
from prefect.utilities.debug import raise_on_exception
with raise_on_exception():
s = f.run(executor=executor)
assert s.result[x].result == [2, 3, 4]
def test(e: Optional[Executor]):
with TemporaryDirectory() as tmpdir:
flow_result = LocalResult(tmpdir, serializer=JSONSerializer(),
location="{task_name}.json")
with Flow("write_result", result=flow_result) as f:
_terminal = task(lambda: 42, checkpoint=True, name="magic")()
with set_temporary_config({"flows.checkpointing": True}), \
raise_on_exception():
f.run(executor=e)
files = os.listdir(tmpdir)
assert files == ["magic.json"], files
with open(os.path.join(tmpdir, files[0]), "rb") as file:
val = json.load(file)
assert val==42
def test_auto_manage_dataframe_default_string(local_file):
task = TaskWithInputDataFrame()
filename = local_file.file.resolve()
with Flow('test_auto_manage_dataframe_default_string') as flow:
file = prefect.Parameter('local_file', default=filename)
task(input_one=file)
with raise_on_exception(), prefect.context(caches={}):
flow_state = flow.run()
result = list(flow_state.result.values())[0].result
df_foo = pd.read_hdf(local_file.file, '/foo')
assert isinstance(df_foo, pd.DataFrame)
tm.assert_equal(result, df_foo.mean())
