Adjust tests

distributed/protocol/pickle.py

@@ -90,6 +90,8 @@ def loads(x, *, buffers=()):
             return pickle.loads(x, buffers=buffers)
         else:
             return pickle.loads(x)
+    except EOFError:
+        raise
     except Exception as e:
         logger.info("Failed to deserialize %s", x[:10000], exc_info=True)
         raise pickle.UnpicklingError("Failed to deserialize") from e

distributed/protocol/tests/test_pickle.py

@@ -15,7 +15,7 @@
 from distributed.protocol import deserialize, serialize
 from distributed.protocol.pickle import HIGHEST_PROTOCOL, dumps, loads
 from distributed.protocol.serialize import dask_deserialize, dask_serialize
-from distributed.utils_test import popen, save_sys_modules
+from distributed.utils_test import popen, raises_with_cause, save_sys_modules
 
 
 class MemoryviewHolder:
@@ -231,7 +231,7 @@ def _deserialize_nopickle(header, frames):
 
 
 def test_allow_pickle_if_registered_in_dask_serialize():
-    with pytest.raises(TypeError, match="nope"):
+    with raises_with_cause(pickle.PicklingError, "serialize", TypeError, "nope"):
         dumps(NoPickle())
 
     dask_serialize.register(NoPickle)(_serialize_nopickle)
@@ -251,9 +251,9 @@ def __init__(self) -> None:
 
 def test_nopickle_nested():
     nested_obj = [NoPickle()]
-    with pytest.raises(TypeError, match="nope"):
+    with raises_with_cause(pickle.PicklingError, "serialize", TypeError, "nope"):
         dumps(nested_obj)
-    with pytest.raises(TypeError, match="nope"):
+    with raises_with_cause(pickle.PicklingError, "serialize", TypeError, "nope"):
         dumps(NestedNoPickle())
 
     dask_serialize.register(NoPickle)(_serialize_nopickle)

distributed/tests/test_client.py

@@ -5100,7 +5100,14 @@ def __setstate__(self, state):
     future = c.submit(identity, Foo())
     await wait(future)
     assert future.status == "error"
-    with raises_with_cause(RuntimeError, "deserialization", MyException, "hello"):
+    with raises_with_cause(
+        RuntimeError,
+        "deserialization",
+        pickle.UnpicklingError,
+        "deserialize",
+        MyException,
+        "hello",
+    ):
         await future
 
     futures = c.map(inc, range(10))
@@ -5125,7 +5132,14 @@ def __call__(self, *args):
     future = c.submit(Foo(), 1)
     await wait(future)
     assert future.status == "error"
-    with raises_with_cause(RuntimeError, "deserialization", MyException, "hello"):
+    with raises_with_cause(
+        RuntimeError,
+        "deserialization",
+        pickle.UnpicklingError,
+        "deserialize",
+        MyException,
+        "hello",
+    ):
         await future
 
     futures = c.map(inc, range(10))

distributed/tests/test_spill.py

@@ -222,7 +222,7 @@ def test_spillbuffer_fail_to_serialize(tmp_path):
     with pytest.raises(TypeError, match="Failed to pickle 'a'") as e:
         with captured_logger("distributed.spill") as logs_bad_key:
             buf["a"] = a
-    assert isinstance(e.value.__cause__.__cause__, MyError)
+    assert isinstance(e.value.__cause__.__cause__.__cause__, MyError)
 
     # spill.py must remain silent because we're already logging in worker.py
     assert not logs_bad_key.getvalue()

distributed/tests/test_worker.py

@@ -6,6 +6,7 @@
 import itertools
 import logging
 import os
+import pickle
 import random
 import sys
 import tempfile
@@ -38,6 +39,7 @@
     get_client,
     get_worker,
     profile,
+    protocol,
     wait,
 )
 from distributed.comm.registry import backends
@@ -46,7 +48,6 @@
 from distributed.core import CommClosedError, Status, rpc
 from distributed.diagnostics.plugin import ForwardOutput
 from distributed.metrics import time
-from distributed.protocol import pickle
 from distributed.scheduler import KilledWorker, Scheduler
 from distributed.utils import get_mp_context, wait_for
 from distributed.utils_test import (
@@ -509,13 +510,15 @@ async def test_plugin_internal_exception():
         with raises_with_cause(
             RuntimeError,
             "Worker failed to start",
+            pickle.UnpicklingError,
+            "deserialize",
             UnicodeDecodeError,
-            match_cause="codec can't decode",
+            "codec can't decode",
         ):
             async with Worker(
                 s.address,
                 plugins={
-                    b"corrupting pickle" + pickle.dumps(lambda: None),
+                    b"corrupting pickle" + protocol.pickle.dumps(lambda: None),
                 },
             ) as w:
                 pass

distributed/tests/test_worker_memory.py

@@ -176,7 +176,7 @@ async def test_fail_to_pickle_execute_1(c, s, a, b):
 
     with pytest.raises(TypeError, match="Failed to pickle 'x'") as e:
         await x
-    assert isinstance(e.value.__cause__.__cause__, CustomError)
+    assert isinstance(e.value.__cause__.__cause__.__cause__, CustomError)
 
     await assert_basic_futures(c)