Fix tests

tests/physical_plan/test_physical_plan_buffering.py

@@ -29,8 +29,11 @@ def test_single_non_buffered_plan():
 
     # Manually "complete" the tasks
     task1.set_result([result1])
+    task1.set_done()
     task2.set_result([result2])
+    task1.set_done()
     task3.set_result([result3])
+    task1.set_done()
 
     # Results should be as we expect
     assert next(plan) == result1
@@ -126,6 +129,7 @@ def test_single_plan_with_buffer_saturation_fast_tasks():
 
     # Finish up on task 1 (task is "fast" and completes so quickly even before the next plan loop call)
     task1.set_result([result1])
+    task1.set_done()
 
     # Plan should fill its buffer completely with new tasks before starting to yield results again
     task2 = next(plan)
@@ -139,7 +143,9 @@ def test_single_plan_with_buffer_saturation_fast_tasks():
 
     # Finish the last task(s)
     task2.set_result([result2])
+    task2.set_done()
     task3.set_result([result3])
+    task3.set_done()
     assert next(plan) == result2
     assert next(plan) == result3