Repository: ray-project/ray
File: python/ray/data/tests/datasource/test_file_based_datasource.py (Python)

Code before:
        ),
    ],
)
def test_read_single_file(ray_start_regular_shared, filesystem, dir_path, endpoint_url):
    # `FileBasedDatasource` should read from the local filesystem if you don't specify
    # one.
    write_filesystem = filesystem
    if write_filesystem is None:
        write_filesystem = pyarrow.fs.LocalFileSystem()

    file_uri = os.path.join(dir_path, "file.txt")

    # PyArrow filesystems expect paths without schemes. `FileBasedDatasource` handles
    # this internally, but we need to manually strip the scheme for the test setup.
    write_path = strip_scheme(file_uri)
    with write_filesystem.open_output_stream(write_path) as f:
        f.write(b"spam")

    datasource = MockFileBasedDatasource(file_uri, filesystem=filesystem)
    tasks = datasource.get_read_tasks(1)
    rows = execute_read_tasks(tasks)

    assert rows == [{"data": b"spam"}]


def test_read_single_directory_local(ray_start_regular_shared, tmp_path):
    dir_path = os.path.join(tmp_path, "dir")
    os.mkdir(dir_path)

    p1 = os.path.join(dir_path, "a.txt")
    with open(p1, "wb") as f:
        f.write(b"a")

    p2 = os.path.join(dir_path, "b.txt")
    with open(p2, "wb") as f:
        f.write(b"b")

    datasource = MockFileBasedDatasource(dir_path)
    rows = execute_read_tasks(datasource.get_read_tasks(1))

    assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"b"}]


def test_read_dir_and_file_mixed(ray_start_regular_shared, tmp_path):
    dir_path = os.path.join(tmp_path, "dir")
    os.mkdir(dir_path)

    p1 = os.path.join(dir_path, "a.txt")
    with open(p1, "wb") as f:
        f.write(b"a")


Review comment (suggestion):
Nit: We don't use the suffix "local" for the other tests even though those use a local filesystem too
```suggestion
def test_read_single_directory(ray_start_regular_shared, tmp_path):
```

Code after:
        ),
    ],
)
def test_read_single_file(ray_start_regular_shared, filesystem, dir_path, endpoint_url):
    # `FileBasedDatasource` should read from the local filesystem if you don't specify
    # one.
    write_filesystem = filesystem
    if write_filesystem is None:
        write_filesystem = pyarrow.fs.LocalFileSystem()

    file_uri = os.path.join(dir_path, "file.txt")

    # PyArrow filesystems expect paths without schemes. `FileBasedDatasource` handles
    # this internally, but we need to manually strip the scheme for the test setup.
    write_path = strip_scheme(file_uri)
    with write_filesystem.open_output_stream(write_path) as f:
        f.write(b"spam")

    datasource = MockFileBasedDatasource(file_uri, filesystem=filesystem)
    tasks = datasource.get_read_tasks(1)
    rows = execute_read_tasks(tasks)

    assert rows == [{"data": b"spam"}]


def test_read_single_directory(ray_start_regular_shared, tmp_path):
    dir_path = tmp_path / "dir"
    dir_path.mkdir()

    p1 = dir_path / "a.txt"
    p1.write_bytes(b"a")

    p2 = dir_path / "b.txt"
    p2.write_bytes(b"b")

    datasource = MockFileBasedDatasource(dir_path)
    rows = execute_read_tasks(datasource.get_read_tasks(1))

    assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"b"}]


def test_read_dir_and_file_mixed(ray_start_regular_shared, tmp_path):
    dir_path = tmp_path / "dir"
    dir_path.mkdir()

    p1 = dir_path / "a.txt"
    p1.write_bytes(b"a")

    p2 = tmp_path / "c.txt"
    p2.write_bytes(b"c")