Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
if self.call_count < 2:
self.call_count += 1
# add small delta to trigger "naptime"
return [pendulum.now("utc").add(seconds=0.05)]
else:
raise SyntaxError("Cease scheduling!")
class StatefulTask(Task):
call_count = 0
def run(self):
self.call_count += 1
t = StatefulTask()
schedule = MockSchedule()
f = Flow(name="test", tasks=[t], schedule=schedule)
with pytest.raises(SyntaxError, match="Cease"):
f.run()
assert t.call_count == 2
def test_sorted_tasks_with_ambiguous_sort():
"""
t1 -> bottleneck
t2 -> bottleneck
t3 -> bottleneck
bottleneck -> t4
bottleneck -> t5
bottleneck -> t6
"""
f = Flow(name="test")
t1 = Task("1")
t2 = Task("2")
t3 = Task("3")
t4 = Task("4")
t5 = Task("5")
t6 = Task("6")
bottleneck = Task("bottleneck")
f.add_edge(t1, bottleneck)
f.add_edge(t2, bottleneck)
f.add_edge(t3, bottleneck)
f.add_edge(bottleneck, t4)
f.add_edge(bottleneck, t5)
f.add_edge(bottleneck, t6)
tasks = f.sorted_tasks()
assert set(tasks[:3]) == set([t1, t2, t3])
def test_load_accepts_name_and_sluggified_name(self):
t = Task(name="foo")
f = Flow("I aM a-test!", tasks=[t])
with tempfile.TemporaryDirectory() as tmpdir:
with set_temporary_config({"home_dir": tmpdir}):
f.save()
new_obj_from_name = Flow.load("I aM a-test!")
new_obj_from_slug = Flow.load("i-am-a-test")
assert isinstance(new_obj_from_name, Flow)
assert len(new_obj_from_name.tasks) == 1
assert list(new_obj_from_name.tasks)[0].name == "foo"
assert list(new_obj_from_name.tasks)[0].slug == t.slug
assert new_obj_from_name.name == "I aM a-test!"
assert isinstance(new_obj_from_slug, Flow)
assert len(new_obj_from_slug.tasks) == 1
assert list(new_obj_from_slug.tasks)[0].name == "foo"
assert list(new_obj_from_slug.tasks)[0].slug == t.slug
assert new_obj_from_slug.name == "I aM a-test!"
def test_update_with_constants():
with Flow("math") as f:
x = Parameter("x")
d = x["d"] + 4
new_flow = Flow("test")
new_flow.update(f)
flow_state = new_flow.run(x=dict(d=42))
assert flow_state.is_successful()
assert flow_state.result[d].result == 46
def test_binding_a_task_in_context_adds_it_to_flow():
with Flow() as flow:
t = Task()
assert t not in flow.tasks
t.bind()
assert t in flow.tasks
def test_flow_dot_run_persists_scheduled_start_time_across_retries(self):
# start very soon
start_time = pendulum.now().add(seconds=0.2)
@task(max_retries=1, retry_delay=datetime.timedelta(0))
def report_start_time():
if prefect.context.task_run_count == 1:
raise ValueError("I'm not ready to tell you the start time yet")
return prefect.context.scheduled_start_time
f = Flow(
name="test",
tasks=[report_start_time],
schedule=prefect.schedules.Schedule(
clocks=[prefect.schedules.clocks.DatesClock(dates=[start_time])]
),
)
state = f.run()
assert state.result[report_start_time].result is start_time
def test_context_manager_is_properly_applied_to_tasks():
t1 = Task()
t2 = Task()
t3 = Task()
with Flow(name="test") as f1:
with Flow(name="test") as f2:
t2.bind()
t1.bind()
with pytest.raises(ValueError):
t3.bind()
assert f1.tasks == set([t1])
assert f2.tasks == set([t2])
def test_flow_parameters():
f = Flow(name="test")
x = Parameter("x")
y = Parameter("y", default=1)
f.add_task(x)
f.add_task(y)
assert f.parameters() == {x, y}
def test_create_flow_without_state_handler(self):
assert Flow().state_handlers == []
def test_copy_creates_new_id():
f = Flow()
f2 = f.copy()
assert f.id != f2.id