add test

ray-project · Mar 3, 2023 · d0c6a1f · d0c6a1f
1 parent ac6f524
commit d0c6a1f
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/python/ray/data/tests/test_util.py b/python/ray/data/tests/test_util.py
@@ -2,7 +2,7 @@
 import ray
 import numpy as np
 
-from ray.data._internal.util import _check_pyarrow_version
+from ray.data._internal.util import _check_pyarrow_version, _split_list
 from ray.data._internal.memory_tracing import (
     trace_allocation,
     trace_deallocation,
@@ -72,6 +72,18 @@ def test_memory_tracing(enabled):
         assert "test5" not in report, report
 
 
+def test_list_splits():
+    with pytest.raises(AssertionError):
+        _split_list(list(range(5)), 0)
+
+    with pytest.raises(AssertionError):
+        _split_list(list(range(5)), -1)
+
+    assert _split_list(list(range(5)), 7) == [[0], [1], [2], [3], [4], [], []]
+    assert _split_list(list(range(5)), 2) == [[0, 1, 2], [3, 4]]
+    assert _split_list(list(range(6)), 2) == [[0, 1, 2], [3, 4, 5]]
+
+
 if __name__ == "__main__":
     import sys