68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
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__]))
|