def test_timeout_handler_reraises():
def do_something():
raise ValueError("test")
with pytest.raises(ValueError) as exc:
timeout_handler(do_something, timeout=1)
assert "test" in exc
def test_timeout_handler_allows_function_to_spawn_new_thread():
def my_thread():
t = threading.Thread(target=lambda: 5)
t.start()
t.join()
assert timeout_handler(my_thread, timeout=1) is None
def test_timeout_handler_allows_function_to_spawn_new_process():
def my_process():
p = multiprocessing.Process(target=lambda: 5)
p.start()
p.join()
p.terminate()
assert timeout_handler(my_process, timeout=1) is None
def test_timeout_handler_preserves_context():
def my_fun(x, **kwargs):
return prefect.context.get("test_key")
with prefect.context(test_key=42):
res = timeout_handler(my_fun, 2, timeout=1)
assert res == 42
def test_timeout_handler_actually_stops_execution():
with tempfile.TemporaryDirectory() as call_dir:
FILE = os.path.join(call_dir, "test.txt")
def slow_fn():
"Runs for 1.5 seconds, writes to file 6 times"
iters = 0
while iters < 6:
time.sleep(0.26)
with open(FILE, "a") as f:
f.write("called\n")
iters += 1
with pytest.raises(TimeoutError):
# allow for at most 3 writes
timeout_handler(slow_fn, timeout=1)
time.sleep(0.5)
with open(FILE, "r") as g:
contents = g.read()
assert len(contents.split("\n")) <= 4
def test_timeout_handler_doesnt_swallow_bad_args():
def do_nothing(x, y=None):
return x, y
with pytest.raises(TypeError):
timeout_handler(do_nothing, timeout=1)
with pytest.raises(TypeError):
timeout_handler(do_nothing, 5, timeout=1, z=10)
with pytest.raises(TypeError):
timeout_handler(do_nothing, 5, timeout=1, y="s", z=10)
def test_timeout_handler_reraises():
def do_something():
raise ValueError("test")
with pytest.raises(ValueError, match="test"):
timeout_handler(do_something, timeout=1)
def test_timeout_handler_passes_args_and_kwargs_and_returns():
def do_nothing(x, y=None):
return x, y
assert timeout_handler(do_nothing, 5, timeout=1, y="yellow") == (5, "yellow")
def test_timeout_handler_times_out():
slow_fn = lambda: time.sleep(2)
with pytest.raises(TimeoutError):
timeout_handler(slow_fn, timeout=1)
def test_timeout_handler_preserves_logging(caplog):
timeout_handler(prefect.Flow("logs").run, timeout=2)
assert len(caplog.records) >= 2 # 1 INFO to start, 1 INFO to end
def test_timeout_handler_doesnt_do_anything_if_no_timeout(monkeypatch):
assert timeout_handler(lambda: 4, timeout=1) == 4
assert timeout_handler(lambda: 4) == 4
def test_timeout_handler_doesnt_do_anything_if_no_timeout(monkeypatch):
monkeypatch.delattr(prefect.utilities.executors, "ThreadPoolExecutor")
with pytest.raises(NameError): # to test the test's usefulness...
timeout_handler(lambda: 4, timeout=1)
assert timeout_handler(lambda: 4) == 4