diff --git a/distributed/tests/test_preload.py b/distributed/tests/test_preload.py index 7a4185f7621..04bec1cb951 100644 --- a/distributed/tests/test_preload.py +++ b/distributed/tests/test_preload.py @@ -1,3 +1,4 @@ +import logging import os import re import shutil @@ -281,6 +282,23 @@ def dask_setup(client, value): assert c.foo == value +@gen_test() +async def test_teardown_failure_doesnt_crash_scheduler(): + text = """ +def dask_teardown(worker): + raise Exception(123) +""" + + with captured_logger(logging.getLogger("distributed.scheduler")) as s_logger: + with captured_logger(logging.getLogger("distributed.worker")) as w_logger: + async with Scheduler(dashboard_address=":0", preload=text) as s: + async with Worker(s.address, preload=[text]) as w: + pass + + assert "123" in s_logger.getvalue() + assert "123" in w_logger.getvalue() + + @gen_cluster(nthreads=[]) async def test_client_preload_config_click(s): text = dedent( @@ -302,14 +320,3 @@ def dask_setup(client, value): ): async with Client(address=s.address, asynchronous=True) as c: assert c.foo == value - - -@gen_test() -async def test_teardown_failure_doesnt_crash_scheduler(): - text = """ -def dask_teardown(worker): - raise Exception(123) -""" - async with Scheduler(dashboard_address=":0", preload=text) as s: - async with Worker(s.address, preload=[text]) as w: - pass