diff --git a/continuous_integration/environment-3.10.yaml b/continuous_integration/environment-3.10.yaml index f48a2df915d..28ccbe31c87 100644 --- a/continuous_integration/environment-3.10.yaml +++ b/continuous_integration/environment-3.10.yaml @@ -47,4 +47,3 @@ dependencies: - git+https://github.com/fsspec/filesystem_spec - git+https://github.com/joblib/joblib - keras - - pytest-asyncio<0.14.0 # `pytest-asyncio<0.14.0` isn't available on conda-forge for Python 3.10 diff --git a/continuous_integration/environment-3.8.yaml b/continuous_integration/environment-3.8.yaml index 715c95a90c3..aba501fb2f1 100644 --- a/continuous_integration/environment-3.8.yaml +++ b/continuous_integration/environment-3.8.yaml @@ -26,7 +26,6 @@ dependencies: - prometheus_client - psutil - pytest - - pytest-asyncio<0.14.0 - pytest-cov - pytest-faulthandler - pytest-repeat diff --git a/continuous_integration/environment-3.9.yaml b/continuous_integration/environment-3.9.yaml index acbb4626b10..8813552c9d7 100644 --- a/continuous_integration/environment-3.9.yaml +++ b/continuous_integration/environment-3.9.yaml @@ -28,7 +28,6 @@ dependencies: - psutil - pynvml # Only tested here - pytest - - pytest-asyncio<0.14.0 - pytest-cov - pytest-faulthandler - pytest-repeat diff --git a/distributed/cli/tests/test_dask_spec.py b/distributed/cli/tests/test_dask_spec.py index eb1a3716da9..d415df9f3a5 100644 --- a/distributed/cli/tests/test_dask_spec.py +++ b/distributed/cli/tests/test_dask_spec.py @@ -39,7 +39,6 @@ async def test_text(): @pytest.mark.skipif(COMPILED, reason="Fails with cythonized scheduler") -@pytest.mark.asyncio @gen_cluster(client=True, nthreads=[]) async def test_file(c, s, tmp_path): fn = str(tmp_path / "foo.yaml") diff --git a/distributed/comm/tests/test_comms.py b/distributed/comm/tests/test_comms.py index f9ecf5e0726..fe1f4855f5f 100644 --- a/distributed/comm/tests/test_comms.py +++ b/distributed/comm/tests/test_comms.py @@ -29,6 +29,7 @@ from distributed.protocol import Serialized, deserialize, serialize, to_serialize from distributed.utils import get_ip, get_ipv6, mp_context from distributed.utils_test import ( + gen_test, get_cert, get_client_ssl_context, get_server_ssl_context, @@ -210,7 +211,7 @@ def test_get_local_address_for(tcp): # -@pytest.mark.asyncio +@gen_test() async def test_tcp_listener_does_not_call_handler_on_handshake_error(tcp): handle_comm_called = False @@ -232,7 +233,7 @@ async def handle_comm(comm): await writer.wait_closed() -@pytest.mark.asyncio +@gen_test() async def test_tcp_specific(tcp): """ Test concrete TCP API. @@ -275,7 +276,7 @@ async def client_communicate(key, delay=0): assert set(l) == {1234} | set(range(N)) -@pytest.mark.asyncio +@gen_test() async def test_tls_specific(tcp): """ Test concrete TLS API. @@ -321,7 +322,7 @@ async def client_communicate(key, delay=0): assert set(l) == {1234} | set(range(N)) -@pytest.mark.asyncio +@gen_test() async def test_comm_failure_threading(tcp): """ When we fail to connect, make sure we don't make a lot @@ -443,17 +444,17 @@ def run(): return fut -@pytest.mark.asyncio +@gen_test() async def test_inproc_specific_same_thread(): await check_inproc_specific(run_coro) -@pytest.mark.asyncio +@gen_test() async def test_inproc_specific_different_threads(): await check_inproc_specific(run_coro_in_thread) -@pytest.mark.asyncio +@gen_test() async def test_inproc_continues_listening_after_handshake_error(): async def handle_comm(): pass @@ -470,7 +471,7 @@ async def handle_comm(): await comm.close() -@pytest.mark.asyncio +@gen_test() async def test_inproc_handshakes_concurrently(): async def handle_comm(): pass @@ -567,7 +568,7 @@ async def client_communicate(key, delay=0): @pytest.mark.gpu -@pytest.mark.asyncio +@gen_test() async def test_ucx_client_server(): pytest.importorskip("distributed.comm.ucx") ucp = pytest.importorskip("ucp") @@ -603,7 +604,7 @@ def checker(loc): return checker -@pytest.mark.asyncio +@gen_test() async def test_default_client_server_ipv4(tcp): # Default scheme is (currently) TCP await check_client_server("127.0.0.1", tcp_eq("127.0.0.1")) @@ -620,7 +621,7 @@ async def test_default_client_server_ipv4(tcp): @requires_ipv6 -@pytest.mark.asyncio +@gen_test() async def test_default_client_server_ipv6(tcp): await check_client_server("[::1]", tcp_eq("::1")) await check_client_server("[::1]:3211", tcp_eq("::1", 3211)) @@ -630,7 +631,7 @@ async def test_default_client_server_ipv6(tcp): ) -@pytest.mark.asyncio +@gen_test() async def test_tcp_client_server_ipv4(tcp): await check_client_server("tcp://127.0.0.1", tcp_eq("127.0.0.1")) await check_client_server("tcp://127.0.0.1:3221", tcp_eq("127.0.0.1", 3221)) @@ -645,7 +646,7 @@ async def test_tcp_client_server_ipv4(tcp): @requires_ipv6 -@pytest.mark.asyncio +@gen_test() async def test_tcp_client_server_ipv6(tcp): await check_client_server("tcp://[::1]", tcp_eq("::1")) await check_client_server("tcp://[::1]:3231", tcp_eq("::1", 3231)) @@ -655,7 +656,7 @@ async def test_tcp_client_server_ipv6(tcp): ) -@pytest.mark.asyncio +@gen_test() async def test_tls_client_server_ipv4(tcp): await check_client_server("tls://127.0.0.1", tls_eq("127.0.0.1"), **tls_kwargs) await check_client_server( @@ -667,12 +668,12 @@ async def test_tls_client_server_ipv4(tcp): @requires_ipv6 -@pytest.mark.asyncio +@gen_test() async def test_tls_client_server_ipv6(tcp): await check_client_server("tls://[::1]", tls_eq("::1"), **tls_kwargs) -@pytest.mark.asyncio +@gen_test() async def test_inproc_client_server(): await check_client_server("inproc://", inproc_check()) await check_client_server(inproc.new_address(), inproc_check()) @@ -683,7 +684,7 @@ async def test_inproc_client_server(): # -@pytest.mark.asyncio +@gen_test() async def test_tls_reject_certificate(tcp): cli_ctx = get_client_ssl_context() serv_ctx = get_server_ssl_context() @@ -755,17 +756,17 @@ async def handle_comm(comm): await comm.read() -@pytest.mark.asyncio +@gen_test() async def test_tcp_comm_closed_implicit(tcp): await check_comm_closed_implicit("tcp://127.0.0.1") -@pytest.mark.asyncio +@gen_test() async def test_tls_comm_closed_implicit(tcp): await check_comm_closed_implicit("tls://127.0.0.1", **tls_kwargs) -@pytest.mark.asyncio +@gen_test() async def test_inproc_comm_closed_implicit(): await check_comm_closed_implicit(inproc.new_address()) @@ -793,22 +794,22 @@ async def check_comm_closed_explicit(addr, listen_args={}, connect_args={}): await b.close() -@pytest.mark.asyncio +@gen_test() async def test_tcp_comm_closed_explicit(tcp): await check_comm_closed_explicit("tcp://127.0.0.1") -@pytest.mark.asyncio +@gen_test() async def test_tls_comm_closed_explicit(tcp): await check_comm_closed_explicit("tls://127.0.0.1", **tls_kwargs) -@pytest.mark.asyncio +@gen_test() async def test_inproc_comm_closed_explicit(): await check_comm_closed_explicit(inproc.new_address()) -@pytest.mark.asyncio +@gen_test() async def test_inproc_comm_closed_explicit_2(): listener_errors = [] @@ -858,7 +859,7 @@ async def handle_comm(comm): await comm.close() -@pytest.mark.asyncio +@gen_test() async def test_comm_closed_on_buffer_error(tcp): # Internal errors from comm.stream.write, such as # BufferError should lead to the stream being closed @@ -890,11 +891,12 @@ async def echo(comm): await comm.write(message) -@pytest.mark.asyncio +@gen_test() async def test_retry_connect(tcp, monkeypatch): async def echo(comm): message = await comm.read() await comm.write(message) + await comm.close() class UnreliableConnector(tcp.TCPConnector): def __init__(self): @@ -921,11 +923,12 @@ class UnreliableBackend(tcp.TCPBackend): await comm.write(b"test") msg = await comm.read() assert msg == b"test" + await comm.close() finally: listener.stop() -@pytest.mark.asyncio +@gen_test() async def test_handshake_slow_comm(tcp, monkeypatch): class SlowComm(tcp.TCP): def __init__(self, *args, delay_in_comm=0.5, **kwargs): @@ -976,12 +979,12 @@ async def check_connect_timeout(addr): assert 1 >= dt >= 0.1 -@pytest.mark.asyncio +@gen_test() async def test_tcp_connect_timeout(tcp): await check_connect_timeout("tcp://127.0.0.1:44444") -@pytest.mark.asyncio +@gen_test() async def test_inproc_connect_timeout(): await check_connect_timeout(inproc.new_address()) @@ -1004,14 +1007,14 @@ async def handle_comm(comm): listener.stop() -@pytest.mark.asyncio +@gen_test() async def test_tcp_many_listeners(tcp): await check_many_listeners("tcp://127.0.0.1") await check_many_listeners("tcp://0.0.0.0") await check_many_listeners("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_inproc_many_listeners(): await check_many_listeners("inproc://") @@ -1160,12 +1163,12 @@ def check_out(deserialize_flag, out_value): await check_connector_deserialize(addr, True, msg, partial(check_out, True)) -@pytest.mark.asyncio +@gen_test() async def test_tcp_deserialize(tcp): await check_deserialize("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_inproc_deserialize(): await check_deserialize("inproc://") @@ -1203,12 +1206,12 @@ async def check_deserialize_roundtrip(addr): assert isinstance(got["ser"], Serialized) -@pytest.mark.asyncio +@gen_test() async def test_inproc_deserialize_roundtrip(): await check_deserialize_roundtrip("inproc://") -@pytest.mark.asyncio +@gen_test() async def test_tcp_deserialize_roundtrip(tcp): await check_deserialize_roundtrip("tcp://") @@ -1238,7 +1241,7 @@ async def handle_comm(comm): await comm.read() -@pytest.mark.asyncio +@gen_test() async def test_tcp_deserialize_eoferror(tcp): await check_deserialize_eoferror("tcp://") @@ -1261,7 +1264,7 @@ async def check_repr(a, b): assert b.peer_address in repr(b) -@pytest.mark.asyncio +@gen_test() async def test_tcp_repr(tcp): a, b = await get_tcp_comm_pair() assert a.local_address in repr(b) @@ -1269,7 +1272,7 @@ async def test_tcp_repr(tcp): await check_repr(a, b) -@pytest.mark.asyncio +@gen_test() async def test_tls_repr(tcp): a, b = await get_tls_comm_pair() assert a.local_address in repr(b) @@ -1277,7 +1280,7 @@ async def test_tls_repr(tcp): await check_repr(a, b) -@pytest.mark.asyncio +@gen_test() async def test_inproc_repr(): a, b = await get_inproc_comm_pair() assert a.local_address in repr(b) @@ -1292,19 +1295,19 @@ async def check_addresses(a, b): b.abort() -@pytest.mark.asyncio +@gen_test() async def test_tcp_adresses(tcp): a, b = await get_tcp_comm_pair() await check_addresses(a, b) -@pytest.mark.asyncio +@gen_test() async def test_tls_adresses(tcp): a, b = await get_tls_comm_pair() await check_addresses(a, b) -@pytest.mark.asyncio +@gen_test() async def test_inproc_adresses(): a, b = await get_inproc_comm_pair() await check_addresses(a, b) diff --git a/distributed/comm/tests/test_ucx.py b/distributed/comm/tests/test_ucx.py index 8e6e9cb4ffe..f4a5729826a 100644 --- a/distributed/comm/tests/test_ucx.py +++ b/distributed/comm/tests/test_ucx.py @@ -14,7 +14,7 @@ from distributed.deploy.local import LocalCluster from distributed.diagnostics.nvml import has_cuda_context from distributed.protocol import to_serialize -from distributed.utils_test import inc +from distributed.utils_test import gen_test, inc try: HOST = ucp.get_address() @@ -61,7 +61,7 @@ async def handle_comm(comm): return (comm, serv_comm) -@pytest.mark.asyncio +@gen_test() async def test_ping_pong(): com, serv_com = await get_comm_pair() msg = {"op": "ping"} @@ -79,7 +79,7 @@ async def test_ping_pong(): await serv_com.close() -@pytest.mark.asyncio +@gen_test() async def test_comm_objs(): comm, serv_comm = await get_comm_pair() @@ -92,7 +92,7 @@ async def test_comm_objs(): assert comm.peer_address == serv_comm.local_address -@pytest.mark.asyncio +@gen_test() async def test_ucx_specific(): """ Test concrete UCX API. @@ -146,7 +146,7 @@ async def client_communicate(key, delay=0): listener.stop() -@pytest.mark.asyncio +@gen_test() async def test_ping_pong_data(): np = pytest.importorskip("numpy") @@ -169,7 +169,7 @@ async def test_ping_pong_data(): await serv_com.close() -@pytest.mark.asyncio +@gen_test() async def test_ucx_deserialize(): # Note we see this error on some systems with this test: # `socket.gaierror: [Errno -5] No address associated with hostname` @@ -179,7 +179,6 @@ async def test_ucx_deserialize(): await check_deserialize("tcp://") -@pytest.mark.asyncio @pytest.mark.parametrize( "g", [ @@ -196,6 +195,7 @@ async def test_ucx_deserialize(): lambda cudf: cudf.DataFrame({"a": ["Check", "str"], "b": ["Sup", "port"]}), ], ) +@gen_test() async def test_ping_pong_cudf(g): # if this test appears after cupy an import error arises # *** ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `CXXABI_1.3.11' @@ -219,8 +219,8 @@ async def test_ping_pong_cudf(g): await serv_com.close() -@pytest.mark.asyncio @pytest.mark.parametrize("shape", [(100,), (10, 10), (4947,)]) +@gen_test() async def test_ping_pong_cupy(shape): cupy = pytest.importorskip("cupy") com, serv_com = await get_comm_pair() @@ -238,8 +238,8 @@ async def test_ping_pong_cupy(shape): @pytest.mark.slow -@pytest.mark.asyncio @pytest.mark.parametrize("n", [int(1e9), int(2.5e9)]) +@gen_test() async def test_large_cupy(n, cleanup): cupy = pytest.importorskip("cupy") com, serv_com = await get_comm_pair() @@ -256,7 +256,7 @@ async def test_large_cupy(n, cleanup): await serv_com.close() -@pytest.mark.asyncio +@gen_test() async def test_ping_pong_numba(): np = pytest.importorskip("numpy") numba = pytest.importorskip("numba") @@ -274,8 +274,8 @@ async def test_ping_pong_numba(): assert result["op"] == "ping" -@pytest.mark.asyncio @pytest.mark.parametrize("processes", [True, False]) +@gen_test() async def test_ucx_localcluster(processes, cleanup): async with LocalCluster( protocol="ucx", @@ -296,7 +296,7 @@ async def test_ucx_localcluster(processes, cleanup): @pytest.mark.slow -@pytest.mark.asyncio +@gen_test(timeout=60) async def test_stress(): da = pytest.importorskip("dask.array") @@ -321,7 +321,7 @@ async def test_stress(): await wait(x) -@pytest.mark.asyncio +@gen_test() async def test_simple(): async with LocalCluster(protocol="ucx", asynchronous=True) as cluster: async with Client(cluster, asynchronous=True) as client: @@ -329,7 +329,7 @@ async def test_simple(): assert await client.submit(lambda x: x + 1, 10) == 11 -@pytest.mark.asyncio +@gen_test() async def test_cuda_context(): with dask.config.set({"distributed.comm.ucx.create-cuda-context": True}): async with LocalCluster( @@ -343,7 +343,7 @@ async def test_cuda_context(): assert list(worker_cuda_context.values())[0] == 0 -@pytest.mark.asyncio +@gen_test() async def test_transpose(): da = pytest.importorskip("dask.array") @@ -356,8 +356,8 @@ async def test_transpose(): await y -@pytest.mark.asyncio @pytest.mark.parametrize("port", [0, 1234]) +@gen_test() async def test_ucx_protocol(cleanup, port): async with Scheduler(protocol="ucx", port=port, dashboard_address=":0") as s: assert s.address.startswith("ucx://") diff --git a/distributed/comm/tests/test_ucx_config.py b/distributed/comm/tests/test_ucx_config.py index 7f41320c229..f1f3f08a3ab 100644 --- a/distributed/comm/tests/test_ucx_config.py +++ b/distributed/comm/tests/test_ucx_config.py @@ -10,7 +10,7 @@ from distributed import Client from distributed.comm.ucx import _scrub_ucx_config from distributed.utils import get_ip -from distributed.utils_test import popen +from distributed.utils_test import gen_test, popen try: HOST = get_ip() @@ -21,7 +21,7 @@ rmm = pytest.importorskip("rmm") -@pytest.mark.asyncio +@gen_test() async def test_ucx_config(cleanup): ucx = { "nvlink": True, diff --git a/distributed/comm/tests/test_ws.py b/distributed/comm/tests/test_ws.py index c75c1a80b7e..acfabaa2bb8 100644 --- a/distributed/comm/tests/test_ws.py +++ b/distributed/comm/tests/test_ws.py @@ -125,7 +125,6 @@ async def test_large_transfer_with_no_compression(): await c.scatter(np.random.random(1_500_000)) -@pytest.mark.asyncio @pytest.mark.parametrize( "dashboard,protocol,security,port", [ @@ -139,6 +138,7 @@ async def test_large_transfer_with_no_compression(): (False, "wss://", True, 8786), ], ) +@gen_test() async def test_http_and_comm_server(cleanup, dashboard, protocol, security, port): if security: xfail_ssl_issue5601() @@ -157,8 +157,8 @@ async def test_http_and_comm_server(cleanup, dashboard, protocol, security, port assert result == 11 -@pytest.mark.asyncio @pytest.mark.parametrize("protocol", ["ws://", "wss://"]) +@gen_test() async def test_connection_made_with_extra_conn_args(cleanup, protocol): if protocol == "ws://": security = Security( diff --git a/distributed/deploy/tests/test_adaptive_core.py b/distributed/deploy/tests/test_adaptive_core.py index 7de0ffe87cc..3e1fa11bee9 100644 --- a/distributed/deploy/tests/test_adaptive_core.py +++ b/distributed/deploy/tests/test_adaptive_core.py @@ -1,10 +1,8 @@ import asyncio -import pytest - from distributed.deploy.adaptive_core import AdaptiveCore from distributed.metrics import time -from distributed.utils_test import captured_logger +from distributed.utils_test import captured_logger, gen_test class MyAdaptive(AdaptiveCore): @@ -25,7 +23,7 @@ async def scale_down(self, workers=()): collection.discard(w) -@pytest.mark.asyncio +@gen_test() async def test_safe_target(): adapt = MyAdaptive(minimum=1, maximum=4) assert await adapt.safe_target() == 1 @@ -33,7 +31,7 @@ async def test_safe_target(): assert await adapt.safe_target() == 4 -@pytest.mark.asyncio +@gen_test() async def test_scale_up(): adapt = MyAdaptive(minimum=1, maximum=4) await adapt.adapt() @@ -46,7 +44,7 @@ async def test_scale_up(): assert adapt.plan == {0, 1, 2, 3} -@pytest.mark.asyncio +@gen_test() async def test_scale_down(): adapt = MyAdaptive(minimum=1, maximum=4, wait_count=2) adapt._target = 10 @@ -72,7 +70,7 @@ async def test_scale_down(): assert list(adapt.log) == old -@pytest.mark.asyncio +@gen_test() async def test_interval(): adapt = MyAdaptive(interval="5 ms") assert not adapt.plan @@ -92,7 +90,7 @@ async def test_interval(): assert len(adapt.plan) == 1 # last value from before, unchanged -@pytest.mark.asyncio +@gen_test() async def test_adapt_oserror_safe_target(): class BadAdaptive(MyAdaptive): """AdaptiveCore subclass which raises an OSError when attempting to adapt @@ -113,7 +111,7 @@ def safe_target(self): assert not adapt.periodic_callback -@pytest.mark.asyncio +@gen_test() async def test_adapt_oserror_scale(): """ FIXME: @@ -147,7 +145,7 @@ async def scale_down(self, workers=None): adapt.stop() -@pytest.mark.asyncio +@gen_test() async def test_adapt_stop_del(): adapt = MyAdaptive(interval="100ms") pc = adapt.periodic_callback diff --git a/distributed/deploy/tests/test_cluster.py b/distributed/deploy/tests/test_cluster.py index 48d8bb65336..b408ee1f435 100644 --- a/distributed/deploy/tests/test_cluster.py +++ b/distributed/deploy/tests/test_cluster.py @@ -1,9 +1,10 @@ import pytest from distributed.deploy.cluster import Cluster +from distributed.utils_test import gen_test -@pytest.mark.asyncio +@gen_test() async def test_eq(cleanup): clusterA = Cluster(asynchronous=True, name="A") clusterA2 = Cluster(asynchronous=True, name="A2") @@ -18,7 +19,7 @@ async def test_eq(cleanup): assert not (clusterA == clusterB) -@pytest.mark.asyncio +@gen_test() async def test_repr(cleanup): cluster = Cluster(asynchronous=True, name="A") assert cluster.scheduler_address == "" @@ -27,14 +28,14 @@ async def test_repr(cleanup): assert res == expected -@pytest.mark.asyncio +@gen_test() async def test_logs_deprecated(cleanup): cluster = Cluster(asynchronous=True) with pytest.warns(FutureWarning, match="get_logs"): cluster.logs() -@pytest.mark.asyncio +@gen_test() async def test_cluster_info(): class FooCluster(Cluster): def __init__(self): diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py index d6ef2298c8e..5c0e755f891 100644 --- a/distributed/deploy/tests/test_local.py +++ b/distributed/deploy/tests/test_local.py @@ -931,8 +931,8 @@ async def test_scale_memory_cores(): assert len(cluster.worker_spec) == 4 -@pytest.mark.asyncio @pytest.mark.parametrize("memory_limit", ["2 GiB", None]) +@gen_test() async def test_repr(memory_limit, cleanup): async with LocalCluster( n_workers=2, @@ -968,8 +968,8 @@ async def test_threads_per_worker_set_to_0(): assert all(w.nthreads < CPU_COUNT for w in cluster.workers.values()) -@pytest.mark.asyncio @pytest.mark.parametrize("temporary", [True, False]) +@gen_test() async def test_capture_security(cleanup, temporary): if temporary: xfail_ssl_issue5601() @@ -1039,28 +1039,29 @@ async def test_cluster_names(): assert unnamed_cluster2 != unnamed_cluster -@pytest.mark.asyncio @pytest.mark.parametrize("nanny", [True, False]) +@gen_test() async def test_local_cluster_redundant_kwarg(nanny): - with pytest.raises(TypeError, match="unexpected keyword argument"): - # Extra arguments are forwarded to the worker class. Depending on - # whether we use the nanny or not, the error treatment is quite - # different and we should assert that an exception is raised - async with await LocalCluster( - typo_kwarg="foo", - processes=nanny, - n_workers=1, - dashboard_address=":0", - ) as cluster: - - # This will never work but is a reliable way to block without hard - # coding any sleep values - async with Client(cluster) as c: - f = c.submit(sleep, 0) - await f + cluster = LocalCluster( + typo_kwarg="foo", + processes=nanny, + n_workers=1, + dashboard_address=":0", + asynchronous=True, + ) + try: + with pytest.raises(TypeError, match="unexpected keyword argument"): + # Extra arguments are forwarded to the worker class. Depending on + # whether we use the nanny or not, the error treatment is quite + # different and we should assert that an exception is raised + async with cluster: + pass + finally: + # FIXME: LocalCluster leaks if LocalCluster.__aenter__ raises + await cluster.close() -@pytest.mark.asyncio +@gen_test() async def test_cluster_info_sync(): async with LocalCluster( processes=False, asynchronous=True, scheduler_sync_interval="1ms" @@ -1087,7 +1088,7 @@ async def test_cluster_info_sync(): assert info["foo"] == "bar" -@pytest.mark.asyncio +@gen_test() async def test_cluster_info_sync_is_robust_to_network_blips(monkeypatch): async with LocalCluster( processes=False, asynchronous=True, scheduler_sync_interval="1ms" @@ -1123,9 +1124,9 @@ async def error(*args, **kwargs): assert info["foo"] == "bar" -@pytest.mark.asyncio @pytest.mark.parametrize("host", [None, "127.0.0.1"]) @pytest.mark.parametrize("use_nanny", [True, False]) +@gen_test() async def test_cluster_host_used_throughout_cluster(host, use_nanny): """Ensure that the `host` kwarg is propagated through scheduler, nanny, and workers""" async with LocalCluster(host=host, asynchronous=True) as cluster: diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py index c6697d1fa75..eea54621a2f 100644 --- a/distributed/deploy/tests/test_slow_adaptive.py +++ b/distributed/deploy/tests/test_slow_adaptive.py @@ -5,7 +5,7 @@ from dask.distributed import Client, Scheduler, SpecCluster, Worker from distributed.metrics import time -from distributed.utils_test import slowinc +from distributed.utils_test import gen_test, slowinc class SlowWorker: @@ -35,7 +35,7 @@ async def close(self): scheduler = {"cls": Scheduler, "options": {"dashboard_address": ":0"}} -@pytest.mark.asyncio +@gen_test() async def test_startup(cleanup): start = time() async with SpecCluster( @@ -55,9 +55,9 @@ async def test_startup(cleanup): await client.wait_for_workers(n_workers=2) -@pytest.mark.asyncio @pytest.mark.flaky(reruns=10, reruns_delay=5) -async def test_scale_up_down(cleanup): +@gen_test() +async def test_scale_up_down(): start = time() async with SpecCluster( scheduler=scheduler, @@ -77,7 +77,7 @@ async def test_scale_up_down(cleanup): assert not cluster.worker_spec -@pytest.mark.asyncio +@gen_test() async def test_adaptive(cleanup): start = time() async with SpecCluster( diff --git a/distributed/deploy/tests/test_spec_cluster.py b/distributed/deploy/tests/test_spec_cluster.py index fe566c94d90..4731e0fe9a2 100644 --- a/distributed/deploy/tests/test_spec_cluster.py +++ b/distributed/deploy/tests/test_spec_cluster.py @@ -21,14 +21,6 @@ class MyWorker(Worker): pass -class BrokenWorker(Worker): - def __await__(self): - async def _(): - raise Exception("Worker Broken") - - return _().__await__() - - worker_spec = { 0: {"cls": "dask.distributed.Worker", "options": {"nthreads": 1}}, 1: {"cls": Worker, "options": {"nthreads": 2}}, @@ -210,20 +202,31 @@ async def test_restart(): @pytest.mark.skipif(WINDOWS, reason="HTTP Server doesn't close out") -# FIXME cleanup fails: -# some RPCs left active by test: {} -# @gen_test() -@pytest.mark.asyncio +@gen_test() async def test_broken_worker(): - with pytest.raises(Exception) as info: - async with SpecCluster( - asynchronous=True, - workers={"good": {"cls": Worker}, "bad": {"cls": BrokenWorker}}, - scheduler=scheduler, - ): - pass + class BrokenWorkerException(Exception): + pass + + class BrokenWorker(Worker): + def __await__(self): + async def _(): + self.status = Status.closed + raise BrokenWorkerException("Worker Broken") + + return _().__await__() - assert "Broken" in str(info.value) + cluster = SpecCluster( + asynchronous=True, + workers={"good": {"cls": Worker}, "bad": {"cls": BrokenWorker}}, + scheduler=scheduler, + ) + try: + with pytest.raises(BrokenWorkerException, match=r"Worker Broken"): + async with cluster: + pass + finally: + # FIXME: SpecCluster leaks if SpecCluster.__aenter__ raises + await cluster.close() @pytest.mark.skipif(WINDOWS, reason="HTTP Server doesn't close out") diff --git a/distributed/deploy/tests/test_ssh.py b/distributed/deploy/tests/test_ssh.py index 61993573ec7..d5d55477941 100644 --- a/distributed/deploy/tests/test_ssh.py +++ b/distributed/deploy/tests/test_ssh.py @@ -9,6 +9,7 @@ from distributed import Client from distributed.compatibility import MACOS, WINDOWS from distributed.deploy.ssh import SSHCluster +from distributed.utils_test import gen_test pytestmark = [ pytest.mark.xfail(MACOS, reason="very high flakiness; see distributed/issues/4543"), @@ -26,7 +27,7 @@ def test_ssh_hosts_empty_list(): SSHCluster(hosts=[]) -@pytest.mark.asyncio +@gen_test() async def test_ssh_cluster_raises_if_asyncssh_not_installed(monkeypatch, cleanup): monkeypatch.setitem(sys.modules, "asyncssh", None) with pytest.raises(ImportError, match="SSHCluster requires the `asyncssh` package"): @@ -40,7 +41,7 @@ async def test_ssh_cluster_raises_if_asyncssh_not_installed(monkeypatch, cleanup assert not cluster -@pytest.mark.asyncio +@gen_test() async def test_basic(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -58,7 +59,7 @@ async def test_basic(): assert "SSH" in repr(cluster) -@pytest.mark.asyncio +@gen_test() async def test_n_workers(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -77,7 +78,7 @@ async def test_n_workers(): assert "SSH" in repr(cluster) -@pytest.mark.asyncio +@gen_test() async def test_nprocs_attribute_is_deprecated(): async with SSHCluster( ["127.0.0.1"] * 2, @@ -97,7 +98,7 @@ async def test_nprocs_attribute_is_deprecated(): assert worker.n_workers == 3 -@pytest.mark.asyncio +@gen_test() async def test_ssh_nprocs_renamed_to_n_workers(): with pytest.warns(FutureWarning, match="renamed to n_workers"): async with SSHCluster( @@ -112,20 +113,25 @@ async def test_ssh_nprocs_renamed_to_n_workers(): await client.wait_for_workers(4) -@pytest.mark.asyncio +@gen_test() async def test_ssh_n_workers_with_nprocs_is_an_error(): - with pytest.raises(ValueError, match="Both nprocs and n_workers"): - async with SSHCluster( - ["127.0.0.1"] * 3, - connect_options=dict(known_hosts=None), - asynchronous=True, - scheduler_options={}, - worker_options={"n_workers": 2, "nprocs": 2}, - ) as cluster: - assert not cluster - - -@pytest.mark.asyncio + cluster = SSHCluster( + ["127.0.0.1"] * 3, + connect_options=dict(known_hosts=None), + asynchronous=True, + scheduler_options={}, + worker_options={"n_workers": 2, "nprocs": 2}, + ) + try: + with pytest.raises(ValueError, match="Both nprocs and n_workers"): + async with cluster: + pass + finally: + # FIXME: SSHCluster leaks if SSHCluster.__aenter__ raises + await cluster.close() + + +@gen_test() async def test_keywords(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -179,7 +185,7 @@ def test_old_ssh_with_local_dir(loop): assert result == 11 -@pytest.mark.asyncio +@gen_test() async def test_config_inherited_by_subprocess(loop): def f(x): return dask.config.get("foo") + 1 @@ -197,7 +203,7 @@ def f(x): assert result == 101 -@pytest.mark.asyncio +@gen_test() async def test_unimplemented_options(): with pytest.raises(Exception): async with SSHCluster( @@ -215,7 +221,7 @@ async def test_unimplemented_options(): assert cluster -@pytest.mark.asyncio +@gen_test() async def test_list_of_connect_options(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -233,7 +239,7 @@ async def test_list_of_connect_options(): assert "SSH" in repr(cluster) -@pytest.mark.asyncio +@gen_test() async def test_list_of_connect_options_raises(): with pytest.raises(RuntimeError): async with SSHCluster( @@ -246,7 +252,7 @@ async def test_list_of_connect_options_raises(): pass -@pytest.mark.asyncio +@gen_test() async def test_remote_python(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -259,7 +265,7 @@ async def test_remote_python(): assert cluster.workers[0].remote_python == sys.executable -@pytest.mark.asyncio +@gen_test() async def test_remote_python_as_dict(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -272,7 +278,7 @@ async def test_remote_python_as_dict(): assert cluster.workers[0].remote_python == sys.executable -@pytest.mark.asyncio +@gen_test() async def test_list_of_remote_python_raises(): with pytest.raises(RuntimeError): async with SSHCluster( diff --git a/distributed/http/tests/test_routing.py b/distributed/http/tests/test_routing.py index 764475ce4ed..9179061d29e 100644 --- a/distributed/http/tests/test_routing.py +++ b/distributed/http/tests/test_routing.py @@ -3,6 +3,7 @@ from tornado.httpclient import AsyncHTTPClient, HTTPClientError from distributed.http.routing import RoutingApplication +from distributed.utils_test import gen_test class OneHandler(web.RequestHandler): @@ -15,7 +16,7 @@ def get(self): self.write("two") -@pytest.mark.asyncio +@gen_test() async def test_basic(): application = RoutingApplication([(r"/one", OneHandler)]) two = web.Application([(r"/two", TwoHandler)]) diff --git a/distributed/protocol/tests/test_serialize.py b/distributed/protocol/tests/test_serialize.py index f144259ba02..8c4eb51c02b 100644 --- a/distributed/protocol/tests/test_serialize.py +++ b/distributed/protocol/tests/test_serialize.py @@ -415,7 +415,7 @@ def dumps(f): assert "Hello-123" in str(info.value) -@pytest.mark.asyncio +@gen_test() async def test_profile_nested_sizeof(): # https://github.com/dask/distributed/issues/1674 n = 500 diff --git a/distributed/tests/test_batched.py b/distributed/tests/test_batched.py index ee84ec3224e..e5cd1037669 100644 --- a/distributed/tests/test_batched.py +++ b/distributed/tests/test_batched.py @@ -9,7 +9,7 @@ from distributed.metrics import time from distributed.protocol import to_serialize from distributed.utils import All -from distributed.utils_test import captured_logger +from distributed.utils_test import captured_logger, gen_test class EchoServer: @@ -37,7 +37,7 @@ async def __aexit__(self, exc, typ, tb): self.stop() -@pytest.mark.asyncio +@gen_test() async def test_BatchedSend(): async with EchoServer() as e: comm = await connect(e.address) @@ -62,9 +62,10 @@ async def test_BatchedSend(): assert result == ("HELLO", "HELLO") assert b.byte_count > 1 + await comm.close() -@pytest.mark.asyncio +@gen_test() async def test_send_before_start(): async with EchoServer() as e: comm = await connect(e.address) @@ -77,9 +78,10 @@ async def test_send_before_start(): b.start(comm) result = await comm.read() assert result == ("hello", "world") + await comm.close() -@pytest.mark.asyncio +@gen_test() async def test_send_after_stream_start(): async with EchoServer() as e: comm = await connect(e.address) @@ -93,9 +95,10 @@ async def test_send_after_stream_start(): if len(result) < 2: result += await comm.read() assert result == ("hello", "world") + await comm.close() -@pytest.mark.asyncio +@gen_test() async def test_send_before_close(): async with EchoServer() as e: comm = await connect(e.address) @@ -117,7 +120,7 @@ async def test_send_before_close(): b.send("123") -@pytest.mark.asyncio +@gen_test() async def test_close_closed(): async with EchoServer() as e: comm = await connect(e.address) @@ -133,13 +136,13 @@ async def test_close_closed(): assert "closed" in str(b) -@pytest.mark.asyncio +@gen_test() async def test_close_not_started(): b = BatchedSend(interval=10) await b.close() -@pytest.mark.asyncio +@gen_test() async def test_close_twice(): async with EchoServer() as e: comm = await connect(e.address) @@ -151,7 +154,7 @@ async def test_close_twice(): @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() async def test_stress(): async with EchoServer() as e: comm = await connect(e.address) @@ -215,18 +218,18 @@ async def run_traffic_jam(nsends, nbytes): await b.close() -@pytest.mark.asyncio +@gen_test() async def test_sending_traffic_jam(): await run_traffic_jam(50, 300000) @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() async def test_large_traffic_jam(): await run_traffic_jam(500, 1500000) -@pytest.mark.asyncio +@gen_test() async def test_serializers(): async with EchoServer() as e: comm = await connect(e.address) diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py index 5044c487bba..2dd848fc80a 100644 --- a/distributed/tests/test_client.py +++ b/distributed/tests/test_client.py @@ -6127,8 +6127,8 @@ async def test_wait_for_workers(c, s, a, b): @pytest.mark.skipif(WINDOWS, reason="num_fds not supported on windows") -@pytest.mark.asyncio @pytest.mark.parametrize("Worker", [Worker, Nanny]) +@gen_test() async def test_file_descriptors_dont_leak(Worker): pytest.importorskip("pandas") df = dask.datasets.timeseries(freq="10s", dtypes={"x": int, "y": float}) @@ -6173,7 +6173,7 @@ async def test_shutdown(): assert w.status == Status.closed -@pytest.mark.asyncio +@gen_test() async def test_shutdown_localcluster(cleanup): async with LocalCluster( n_workers=1, asynchronous=True, processes=False, dashboard_address=":0" @@ -7402,7 +7402,7 @@ def config_loader(self, monkeypatch, loader): ): yield - @pytest.mark.asyncio + @gen_test() async def test_security_loader(self, monkeypatch): security = tls_only_security() @@ -7418,7 +7418,7 @@ def loader(info): async with Client(scheduler.address, asynchronous=True) as client: assert client.security is security - @pytest.mark.asyncio + @gen_test() async def test_security_loader_ignored_if_explicit_security_provided( self, monkeypatch ): @@ -7436,7 +7436,7 @@ def loader(info): ) as client: assert client.security is security - @pytest.mark.asyncio + @gen_test() async def test_security_loader_ignored_if_returns_none(self, monkeypatch): """Test that if a security loader is configured, but it returns `None`, then the default security configuration is used""" @@ -7469,7 +7469,7 @@ def loader(info): assert loader.called - @pytest.mark.asyncio + @gen_test() async def test_security_loader_import_failed(self): security = tls_only_security() diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py index 3598a00bbc9..a32b7432929 100644 --- a/distributed/tests/test_core.py +++ b/distributed/tests/test_core.py @@ -1,10 +1,12 @@ import asyncio +import contextlib import os import socket import threading import weakref import pytest +from tornado.ioloop import IOLoop import dask @@ -145,7 +147,7 @@ class MyServer(Server): default_port = 8756 -@pytest.mark.asyncio +@gen_test() async def test_server_listen(): """ Test various Server.listen() arguments and their effect. @@ -291,7 +293,7 @@ async def check_rpc(listen_addr, rpc_addr=None, listen_args={}, connection_args= await asyncio.sleep(0) -@pytest.mark.asyncio +@gen_test() async def test_rpc_default(): await check_rpc(8883, "127.0.0.1:8883") await check_rpc(8883) @@ -303,7 +305,7 @@ async def test_rpc_tcp(): await check_rpc("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_rpc_tls(): sec = tls_security() await check_rpc( @@ -314,12 +316,12 @@ async def test_rpc_tls(): ) -@pytest.mark.asyncio +@gen_test() async def test_rpc_inproc(): await check_rpc("inproc://", None) -@pytest.mark.asyncio +@gen_test() async def test_rpc_inputs(): L = [rpc("127.0.0.1:8884"), rpc(("127.0.0.1", 8884)), rpc("tcp://127.0.0.1:8884")] @@ -363,17 +365,17 @@ async def check_rpc_message_lifetime(*listen_args): server.stop() -@pytest.mark.asyncio +@gen_test() async def test_rpc_message_lifetime_default(): await check_rpc_message_lifetime() -@pytest.mark.asyncio +@gen_test() async def test_rpc_message_lifetime_tcp(): await check_rpc_message_lifetime("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_rpc_message_lifetime_inproc(): await check_rpc_message_lifetime("inproc://") @@ -396,12 +398,12 @@ async def g(): assert all(comm.closed() for comm in remote.comms) -@pytest.mark.asyncio +@gen_test() async def test_rpc_with_many_connections_tcp(): await check_rpc_with_many_connections("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_rpc_with_many_connections_inproc(): await check_rpc_with_many_connections("inproc://") @@ -424,12 +426,12 @@ async def check_large_packets(listen_arg): @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() async def test_large_packets_tcp(): await check_large_packets("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_large_packets_inproc(): await check_large_packets("inproc://") @@ -447,18 +449,19 @@ async def check_identity(listen_arg): server.stop() -@pytest.mark.asyncio +@gen_test() async def test_identity_tcp(): await check_identity("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_identity_inproc(): await check_identity("inproc://") -@pytest.mark.asyncio -async def test_ports(loop): +@gen_test() +async def test_ports(): + loop = IOLoop.current() for port in range(9877, 9887): server = Server({}, io_loop=loop) try: @@ -491,7 +494,7 @@ def stream_div(comm=None, x=None, y=None): return x / y -@pytest.mark.asyncio +@gen_test() async def test_errors(): server = Server({"div": stream_div}) await server.listen(0) @@ -501,13 +504,13 @@ async def test_errors(): await r.div(x=1, y=0) -@pytest.mark.asyncio +@gen_test() async def test_connect_raises(): with pytest.raises(IOError): await connect("127.0.0.1:58259", timeout=0.01) -@pytest.mark.asyncio +@gen_test() async def test_send_recv_args(): server = Server({}) await server.listen(0) @@ -553,7 +556,7 @@ def test_coerce_to_address(): assert coerce_to_address(arg) == "tcp://127.0.0.1:8786" -@pytest.mark.asyncio +@gen_test() async def test_connection_pool(): async def ping(comm, delay=0.1): await asyncio.sleep(delay) @@ -602,7 +605,7 @@ async def ping(comm, delay=0.1): await rpc.close() -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_close_while_connecting(monkeypatch): """ Ensure a closed connection pool guarantees to have no connections left open @@ -646,7 +649,7 @@ async def connect_to_server(): assert not pool._n_connecting -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_outside_cancellation(monkeypatch): # Ensure cancellation errors are properly reraised from distributed.comm.registry import backends @@ -684,7 +687,7 @@ async def connect_to_server(): assert all(t.cancelled() for t in tasks) -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_respects_limit(): limit = 5 @@ -698,16 +701,19 @@ async def do_ping(pool, port): await pool(ip="127.0.0.1", port=port).ping() assert pool.open <= limit - servers = [Server({"ping": ping}) for i in range(10)] - for server in servers: - await server.listen(0) - - pool = await ConnectionPool(limit=limit) + async with contextlib.AsyncExitStack() as stack: + servers = [ + await stack.enter_async_context(Server({"ping": ping})) for i in range(10) + ] + for server in servers: + await server.listen(0) - await asyncio.gather(*(do_ping(pool, s.port) for s in servers)) + pool = await ConnectionPool(limit=limit) + await asyncio.gather(*(do_ping(pool, s.port) for s in servers)) + await pool.close() -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_tls(): """ Make sure connection args are supported. @@ -734,7 +740,7 @@ async def ping(comm, delay=0.01): await rpc.close() -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_remove(): async def ping(comm, delay=0.01): await asyncio.sleep(delay) @@ -773,7 +779,7 @@ async def ping(comm, delay=0.01): await rpc.close() -@pytest.mark.asyncio +@gen_test() async def test_counters(): server = Server({"div": stream_div}) await server.listen("tcp://") @@ -832,7 +838,7 @@ async def f(): loop.run_sync(f) -@pytest.mark.asyncio +@gen_test() async def test_rpc_serialization(): server = Server({"echo": echo_serialize}) await server.listen("tcp://") @@ -853,21 +859,23 @@ async def test_thread_id(s, a, b): assert s.thread_id == a.thread_id == b.thread_id == threading.get_ident() -@pytest.mark.asyncio +@gen_test() async def test_deserialize_error(): server = Server({"throws": throws}) await server.listen(0) comm = await connect(server.address, deserialize=False) - with pytest.raises(Exception) as info: + with pytest.raises(Exception, match=r"RuntimeError\('hello!'\)") as info: await send_recv(comm, op="throws", x="foo") assert type(info.value) == Exception for c in str(info.value): assert c.isalpha() or c in "(',!)" # no crazy bytestrings + await comm.close() + -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_detects_remote_close(): server = Server({"ping": pingpong}) await server.listen("tcp://") @@ -899,7 +907,7 @@ async def test_connection_pool_detects_remote_close(): await p.close() -@pytest.mark.asyncio +@gen_test() async def test_close_properly(): """ If the server is closed we should cancel all still ongoing coros and close @@ -947,13 +955,13 @@ async def sleep(comm=None): assert not len(server._ongoing_coroutines) -@pytest.mark.asyncio +@gen_test() async def test_server_redundant_kwarg(): with pytest.raises(TypeError, match="unexpected keyword argument"): await Server({}, typo_kwarg="foo") -@pytest.mark.asyncio +@gen_test() async def test_server_comms_mark_active_handlers(): """Whether handlers are active can be read off of the self._comms values. ensure this is properly reflected and released. The sentinel for @@ -980,7 +988,7 @@ async def long_handler(comm): await asyncio.sleep(0.01) -@pytest.mark.asyncio +@gen_test() async def test_close_fast_without_active_handlers(): async def very_fast(comm): return "done" @@ -998,7 +1006,7 @@ async def very_fast(comm): await asyncio.wait_for(fut, 0.1) -@pytest.mark.asyncio +@gen_test() async def test_close_grace_period_for_handlers(): async def long_handler(comm, delay=10): await asyncio.sleep(delay) diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py index 33cffc8d72a..6312a4dfe9f 100644 --- a/distributed/tests/test_nanny.py +++ b/distributed/tests/test_nanny.py @@ -499,7 +499,7 @@ def raise_err(): @pytest.mark.parametrize("protocol", ["tcp", "ucx"]) -@pytest.mark.asyncio +@gen_test() async def test_nanny_closed_by_keyboard_interrupt(cleanup, protocol): if protocol == "ucx": # Skip if UCX isn't available pytest.importorskip("ucp") diff --git a/distributed/tests/test_preload.py b/distributed/tests/test_preload.py index 6d4c8b1fce8..073c9a2d959 100644 --- a/distributed/tests/test_preload.py +++ b/distributed/tests/test_preload.py @@ -1,26 +1,18 @@ -import multiprocessing import os +import re import shutil import sys import tempfile -import urllib.error -import urllib.request from textwrap import dedent -from time import sleep +from unittest import mock import pytest -import tornado -from tornado import web import dask from distributed import Client, Nanny, Scheduler, Worker -from distributed.compatibility import MACOS -from distributed.metrics import time from distributed.utils_test import captured_logger, cluster, gen_cluster, gen_test -PY_VERSION = sys.version_info[:2] - PRELOAD_TEXT = """ _worker_info = {} @@ -142,7 +134,7 @@ async def dask_setup(worker): assert w.foo == "setup" -@pytest.mark.asyncio +@gen_test() async def test_preload_import_time(cleanup): text = """ from distributed.comm.registry import backends @@ -161,61 +153,32 @@ async def test_preload_import_time(cleanup): del backends["foo"] -class MyHandler(web.RequestHandler): - def get(self): - self.write( - """ -def dask_setup(dask_server): - dask_server.foo = 1 -""".strip() - ) - - -def create_preload_application(): - app = web.Application([(r"/preload", MyHandler)]) - server = app.listen(12345, address="127.0.0.1") - tornado.ioloop.IOLoop.instance().start() - - -@pytest.fixture -def scheduler_preload(): - p = multiprocessing.Process(target=create_preload_application) - p.start() - start = time() - while not p.is_alive(): - if time() > start + 5: - raise AssertionError("Process didn't come up") - sleep(0.5) - # Make sure we can query the server - start = time() - request = urllib.request.Request("http://127.0.0.1:12345/preload", method="GET") - while True: - try: - response = urllib.request.urlopen(request) - if response.status == 200: - break - except urllib.error.URLError as e: - if time() > start + 10: - raise AssertionError("Webserver didn't come up", e) - sleep(0.5) - - yield - p.kill() - p.join(timeout=5) - - -@pytest.mark.skipif( - MACOS and PY_VERSION == (3, 7), reason="HTTP Server doesn't come up" -) -@pytest.mark.asyncio -async def test_web_preload(cleanup, scheduler_preload): - with captured_logger("distributed.preloading") as log: +@gen_test() +async def test_web_preload(): + with mock.patch( + "urllib3.PoolManager.request", + **{ + "return_value.data": b"def dask_setup(dask_server):" + b"\n dask_server.foo = 1" + b"\n" + }, + ) as request, captured_logger("distributed.preloading") as log: async with Scheduler( - host="localhost", - preload=["http://127.0.0.1:12345/preload"], + host="localhost", preload=["http://example.com/preload"] ) as s: assert s.foo == 1 - assert "12345/preload" in log.getvalue() + assert ( + re.match( + r"(?s).*Downloading preload at http://example.com/preload\n" + r".*Run preload setup function: http://example.com/preload\n" + r".*", + log.getvalue(), + ) + is not None + ) + assert request.mock_calls == [ + mock.call(method="GET", url="http://example.com/preload", retries=mock.ANY) + ] @gen_cluster(nthreads=[]) @@ -238,57 +201,22 @@ async def test_scheduler_startup_nanny(s): assert w.scheduler.address == s.address -class WorkerPreloadHandler(web.RequestHandler): - def get(self): - self.write( - """ -import dask -dask.config.set(scheduler_address="tcp://127.0.0.1:8786") -""".strip() - ) - - -def create_worker_preload_application(): - application = web.Application([(r"/preload", WorkerPreloadHandler)]) - server = application.listen(12346, address="127.0.0.1") - tornado.ioloop.IOLoop.instance().start() - - -@pytest.fixture -def worker_preload(): - p = multiprocessing.Process(target=create_worker_preload_application) - p.start() - start = time() - while not p.is_alive(): - if time() > start + 5: - raise AssertionError("Process didn't come up") - sleep(0.5) - # Make sure we can query the server - request = urllib.request.Request("http://127.0.0.1:12346/preload", method="GET") - start = time() - while True: - try: - response = urllib.request.urlopen(request) - if response.status == 200: - break - except urllib.error.URLError as e: - if time() > start + 10: - raise AssertionError("Webserver didn't come up", e) - sleep(0.5) - - yield - p.kill() - p.join(timeout=5) - - -@pytest.mark.skipif( - MACOS and PY_VERSION == (3, 7), reason="HTTP Server doesn't come up" -) -@pytest.mark.asyncio -async def test_web_preload_worker(cleanup, worker_preload): - async with Scheduler(port=8786, host="localhost") as s: - async with Nanny(preload_nanny=["http://127.0.0.1:12346/preload"]) as nanny: - assert nanny.scheduler_addr == s.address +@gen_test() +async def test_web_preload_worker(): + with mock.patch( + "urllib3.PoolManager.request", + **{ + "return_value.data": b"import dask" + b'\ndask.config.set(scheduler_address="tcp://127.0.0.1:8786")' + b"\n" + }, + ) as request: + async with Scheduler(port=8786, host="localhost") as s: + async with Nanny(preload_nanny=["http://example.com/preload"]) as nanny: + assert nanny.scheduler_addr == s.address + assert request.mock_calls == [ + mock.call(method="GET", url="http://example.com/preload", retries=mock.ANY) + ] # This test is blocked on https://github.com/dask/distributed/issues/5819 diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py index 2750868dc01..7ee7e2fe0d9 100644 --- a/distributed/tests/test_scheduler.py +++ b/distributed/tests/test_scheduler.py @@ -1350,7 +1350,7 @@ async def test_non_existent_worker(c, s): ("127.0.0.1:0", ("127.0.0.1",)), ], ) -@pytest.mark.asyncio +@gen_test() async def test_dashboard_host(host, dashboard_address, expect): """Dashboard is accessible from any host by default, but it can be also bound to localhost. diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py index c920366c4c2..51fd315a26c 100644 --- a/distributed/tests/test_utils.py +++ b/distributed/tests/test_utils.py @@ -563,13 +563,13 @@ def test_lru(): assert list(l.keys()) == ["c", "a", "d"] -@pytest.mark.asyncio +@gen_test() async def test_offload(): assert (await offload(inc, 1)) == 2 assert (await offload(lambda x, y: x + y, 1, y=2)) == 3 -@pytest.mark.asyncio +@gen_test() async def test_offload_preserves_contextvars(): var = contextvars.ContextVar("var") diff --git a/distributed/tests/test_utils_test.py b/distributed/tests/test_utils_test.py index dc481aae17f..9cdbadf2379 100755 --- a/distributed/tests/test_utils_test.py +++ b/distributed/tests/test_utils_test.py @@ -263,7 +263,7 @@ def test_tls_cluster(tls_client): assert tls_client.security -@pytest.mark.asyncio +@gen_test() async def test_tls_scheduler(security, cleanup): async with Scheduler( security=security, host="localhost", dashboard_address=":0" @@ -289,7 +289,7 @@ def pong(self, comm): return "pong" -@pytest.mark.asyncio +@gen_test() async def test_locked_comm_drop_in_replacement(loop): async with MyServer({}) as a, await MyServer({}) as b: @@ -318,7 +318,7 @@ async def test_locked_comm_drop_in_replacement(loop): assert await read_queue.get() == (b.address, "pong") -@pytest.mark.asyncio +@gen_test() async def test_locked_comm_intercept_read(loop): async with MyServer({}) as a, MyServer({}) as b: @@ -347,7 +347,7 @@ async def ping_pong(): assert await fut == "pong" -@pytest.mark.asyncio +@gen_test() async def test_locked_comm_intercept_write(loop): async with MyServer({}) as a, MyServer({}) as b: diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py index aa8549b07ba..357c34e5603 100644 --- a/distributed/tests/test_worker.py +++ b/distributed/tests/test_worker.py @@ -389,7 +389,7 @@ def __str__(self): assert "Bar" in str(e.__cause__) -@pytest.mark.asyncio +@gen_test() async def test_plugin_exception(cleanup): class MyPlugin: def setup(self, worker=None): @@ -406,7 +406,7 @@ def setup(self, worker=None): pass -@pytest.mark.asyncio +@gen_test() async def test_plugin_multiple_exceptions(cleanup): class MyPlugin1: def setup(self, worker=None): @@ -434,7 +434,7 @@ def setup(self, worker=None): assert "MyPlugin2 Error" in text -@pytest.mark.asyncio +@gen_test() async def test_plugin_internal_exception(cleanup): async with Scheduler(port=0) as s: with pytest.raises(UnicodeDecodeError, match="codec can't decode"): @@ -1323,9 +1323,9 @@ async def test_host_address(c, s): await n.close() -@pytest.mark.asyncio @pytest.mark.parametrize("Worker", [Worker, Nanny]) -async def test_interface_async(cleanup, loop, Worker): +@gen_test() +async def test_interface_async(Worker): from distributed.utils import get_ip_interface psutil = pytest.importorskip("psutil") @@ -1356,9 +1356,9 @@ async def test_interface_async(cleanup, loop, Worker): @pytest.mark.gpu -@pytest.mark.asyncio @pytest.mark.parametrize("Worker", [Worker, Nanny]) -async def test_protocol_from_scheduler_address(cleanup, Worker): +@gen_test() +async def test_protocol_from_scheduler_address(Worker): pytest.importorskip("ucp") async with Scheduler(protocol="ucx", dashboard_address=":0") as s: @@ -1370,7 +1370,7 @@ async def test_protocol_from_scheduler_address(cleanup, Worker): assert info["address"].startswith("ucx://") -@pytest.mark.asyncio +@gen_test() async def test_host_uses_scheduler_protocol(cleanup, monkeypatch): # Ensure worker uses scheduler's protocol to determine host address, not the default scheme # See https://github.com/dask/distributed/pull/4883 @@ -1390,9 +1390,9 @@ def get_address_host(self, loc): pass -@pytest.mark.asyncio @pytest.mark.parametrize("Worker", [Worker, Nanny]) -async def test_worker_listens_on_same_interface_by_default(cleanup, Worker): +@gen_test() +async def test_worker_listens_on_same_interface_by_default(Worker): async with Scheduler(host="localhost", dashboard_address=":0") as s: assert s.ip in {"127.0.0.1", "localhost"} async with Worker(s.address) as w: