fix memory buffer tests

fjetter · fjetter · commit 67eb5457e071 · 2024-02-22T10:53:21.000+01:00
diff --git a/distributed/shuffle/tests/test_memory_buffer.py b/distributed/shuffle/tests/test_memory_buffer.py
@@ -2,6 +2,7 @@
 
 import pytest
 
+from distributed.shuffle._exceptions import DataUnavailable
 from distributed.shuffle._memory import MemoryShardsBuffer
 from distributed.utils_test import gen_test
 
@@ -21,7 +22,7 @@ async def test_basic():
         x = mf.read("x")
         y = mf.read("y")
 
-        with pytest.raises(KeyError):
+        with pytest.raises(DataUnavailable):
             mf.read("z")
 
         assert x == [b"0" * 1000] * 2
@@ -42,7 +43,7 @@ async def test_read_before_flush():
 
         await mf.flush()
         assert mf.read("1") == [b"foo"]
-        with pytest.raises(KeyError):
+        with pytest.raises(DataUnavailable):
             mf.read("2")
 
 
