diff --git a/packages/vaex-core/vaex/test/dataset.py b/packages/vaex-core/vaex/test/dataset.py
index <HASH>..<HASH> 100644
--- a/packages/vaex-core/vaex/test/dataset.py
+++ b/packages/vaex-core/vaex/test/dataset.py
@@ -23,8 +23,10 @@ basedir = os.path.dirname(__file__)
 def small_buffer(ds, size=3):
 	previous = ds.executor.buffer_size
 	ds.executor.buffer_size = size
-	yield
-	ds.executor.buffer_size = previous
+	try:
+		yield
+	finally:
+		ds.executor.buffer_size = previous
 
 
 # these need to be global for pickling
test: fix: on failure, the buffer size wasnt't restored, causing the unittest to slow down