Test get_default_compression

distributed/protocol/tests/test_protocol.py

@@ -1,7 +1,13 @@
 import pytest
 
+import dask
+
 from distributed.protocol import dumps, loads, maybe_compress, msgpack, to_serialize
-from distributed.protocol.compression import compressions
+from distributed.protocol.compression import (
+    compressions,
+    default_compression,
+    get_default_compression,
+)
 from distributed.protocol.cuda import cuda_deserialize, cuda_serialize
 from distributed.protocol.serialize import (
     Serialize,
@@ -20,6 +26,24 @@ def test_protocol():
         assert loads(dumps(msg)) == msg
 
 
+@pytest.mark.parametrize(
+    "config,default",
+    [
+        ("auto", default_compression),
+        (None, None),
+        ("zlib", "zlib"),
+        ("foo", ValueError),
+    ],
+)
+def test_compression_config(config, default):
+    with dask.config.set({"distributed.comm.compression": config}):
+        if type(default) is type and issubclass(default, Exception):
+            with pytest.raises(default):
+                assert get_default_compression()
+        else:
+            assert get_default_compression() == default
+
+
 def test_compression_1():
     pytest.importorskip("lz4")
     np = pytest.importorskip("numpy")