diff --git a/distributed/protocol/tests/test_to_pickle.py b/distributed/protocol/tests/test_to_pickle.py index 7db7a5d9738..d3099c9a73a 100644 --- a/distributed/protocol/tests/test_to_pickle.py +++ b/distributed/protocol/tests/test_to_pickle.py @@ -4,10 +4,22 @@ from dask.highlevelgraph import HighLevelGraph, MaterializedLayer from distributed.client import Client +from distributed.protocol import dumps, loads from distributed.protocol.serialize import ToPickle from distributed.utils_test import gen_cluster +def test_ToPickle(): + class Foo: + def __init__(self, data): + self.data = data + + msg = {"x": ToPickle(Foo(123))} + frames = dumps(msg) + out = loads(frames) + assert out["x"].data == 123 + + class NonMsgPackSerializableLayer(MaterializedLayer): """Layer that uses non-msgpack-serializable data"""