chore: import upstream snapshot with attribution

python/ray/data/tests/unit/test_filename_provider.py

@@ -0,0 +1,67 @@
+import pytest
+
+from ray.data.datasource.filename_provider import FilenameProvider
+
+
+@pytest.fixture(params=["csv", None])
+def filename_provider(request):
+    yield FilenameProvider(dataset_uuid="", file_format=request.param)
+
+
+def test_default_filename_for_task_includes_task_index(filename_provider):
+    filename_0 = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+    filename_1 = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=1
+    )
+    assert filename_0 != filename_1
+    assert "000000" in filename_0
+    assert "000001" in filename_1
+
+
+def test_default_get_filename_for_task_is_deterministic(filename_provider):
+    """Test the new get_filename_for_task() method is deterministic."""
+
+    first_filename = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+    second_filename = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+
+    assert first_filename == second_filename
+
+
+def test_default_row_filenames_derived_from_task_are_unique(filename_provider):
+    """Row filenames derived from task filename with block/row index are unique."""
+    task_filename = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+    filenames = []
+    for block_index in range(2):
+        for row_index in range(4):
+            if "." in task_filename:
+                base, ext = task_filename.rsplit(".", 1)
+                filenames.append(f"{base}_{block_index:06}_{row_index:06}.{ext}")
+            else:
+                filenames.append(f"{task_filename}_{block_index:06}_{row_index:06}")
+    assert len(set(filenames)) == len(filenames)
+
+
+def test_default_get_filename_for_task_is_unique(filename_provider):
+    """Test the new get_filename_for_task() method generates unique filenames."""
+    filenames = [
+        filename_provider.get_filename_for_task(
+            write_uuid="spam",
+            task_index=task_index,
+        )
+        for task_index in range(4)
+    ]
+    assert len(set(filenames)) == len(filenames)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))