Fix flaky test_shuffle_extension.py::test_get_partition
#5689
+9
−7
test_shuffle_extension.py::test_get_partition
#5689