timsaucer commented on code in PR #1625: URL: https://github.com/apache/datafusion-python/pull/1625#discussion_r3683993919
########## python/tests/test_object_store_param.py: ########## @@ -0,0 +1,301 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Tests for the object_store parameter on register/read file methods.""" + +import contextlib +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pyarrow as pa +import pytest +from datafusion import SessionContext + + [email protected] +def ctx(): + return SessionContext() + + +class TestRegisterObjectStoreForPath: + """Unit tests for _register_object_store_for_path URL parsing logic.""" + + def test_parses_s3_url(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + ctx._register_object_store_for_path( + "s3://my-bucket/path/to/file.parquet", mock_store + ) + mock_register.assert_called_once_with("s3://", mock_store, host="my-bucket") + + def test_parses_gs_url(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + ctx._register_object_store_for_path( + "gs://my-gcs-bucket/data.parquet", mock_store + ) + mock_register.assert_called_once_with( + "gs://", mock_store, host="my-gcs-bucket" + ) + + def test_parses_az_url(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + ctx._register_object_store_for_path( + "az://my-container/data.parquet", mock_store + ) + mock_register.assert_called_once_with( + "az://", mock_store, host="my-container" + ) + + def test_parses_https_url(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + ctx._register_object_store_for_path( + "https://my-host.example.com/data.parquet";, mock_store + ) + mock_register.assert_called_once_with( + "https://";, mock_store, host="my-host.example.com" + ) + + def test_raises_on_local_path(self, ctx): + mock_store = MagicMock() + with pytest.raises(ValueError, match="Cannot determine object store URL"): + ctx._register_object_store_for_path("/local/path/file.parquet", mock_store) + + def test_raises_on_relative_path(self, ctx): + mock_store = MagicMock() + with pytest.raises(ValueError, match="Cannot determine object store URL"): + ctx._register_object_store_for_path("relative/path.parquet", mock_store) + + def test_raises_on_windows_path(self, ctx): + mock_store = MagicMock() + with pytest.raises(ValueError, match="must include a host or bucket"): + ctx._register_object_store_for_path( + "C:\\Users\\data\\file.parquet", mock_store + ) + + def test_raises_on_scheme_without_host_for_non_file(self, ctx): + """Non-file schemes (s3, gs, etc.) require a host/bucket.""" + mock_store = MagicMock() + with pytest.raises(ValueError, match="must include a host or bucket"): + ctx._register_object_store_for_path("s3:///key.parquet", mock_store) + + def test_parses_file_url(self, ctx): + """file:// URLs with empty netloc should be accepted.""" + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + ctx._register_object_store_for_path( + "file:///tmp/path/to/file.parquet", mock_store + ) + mock_register.assert_called_once_with("file://", mock_store, host=None) + + def test_accepts_pathlib_path_raises(self, ctx): + """pathlib.Path cannot represent URLs, so this should raise.""" + mock_store = MagicMock() + # pathlib.Path strips the scheme, so this becomes a local path + with pytest.raises(ValueError, match="Cannot determine object store URL"): + ctx._register_object_store_for_path(Path("/local/file.parquet"), mock_store) + + +class TestRegisterParquetObjectStore: + """Tests for register_parquet with object_store parameter.""" + + def test_object_store_none_does_not_register(self, ctx): + """When object_store is None, register_object_store is not called.""" + with patch.object(ctx, "register_object_store") as mock_register: + # This will fail at the Rust level (file doesn't exist), but + # we're testing that register_object_store is NOT called + with contextlib.suppress(Exception): + ctx.register_parquet("t", "s3://bucket/file.parquet") + mock_register.assert_not_called() + + def test_object_store_triggers_registration(self, ctx): + """When object_store is provided, register_object_store is called.""" + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.register_parquet( + "t", + "s3://my-bucket/file.parquet", + object_store=mock_store, + ) + mock_register.assert_called_once_with("s3://", mock_store, host="my-bucket") + + def test_object_store_invalid_path_raises(self, ctx): + """Providing object_store with a local path raises ValueError.""" + mock_store = MagicMock() + with pytest.raises(ValueError, match="Cannot determine object store URL"): + ctx.register_parquet("t", "/local/file.parquet", object_store=mock_store) + + +class TestReadParquetObjectStore: + """Tests for read_parquet with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.read_parquet("s3://my-bucket/file.parquet", object_store=mock_store) + mock_register.assert_called_once_with("s3://", mock_store, host="my-bucket") + + +class TestRegisterCsvObjectStore: + """Tests for register_csv with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.register_csv( + "t", "s3://my-bucket/data.csv", object_store=mock_store + ) + mock_register.assert_called_once_with("s3://", mock_store, host="my-bucket") + + def test_object_store_with_list_path(self, ctx): + """For list paths, the first entry is used for URL parsing.""" + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.register_csv( + "t", + ["s3://my-bucket/a.csv", "s3://my-bucket/b.csv"], + object_store=mock_store, + ) + mock_register.assert_called_once_with("s3://", mock_store, host="my-bucket") + + +class TestReadCsvObjectStore: + """Tests for read_csv with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.read_csv("gs://bucket/data.csv", object_store=mock_store) + mock_register.assert_called_once_with("gs://", mock_store, host="bucket") + + +class TestRegisterJsonObjectStore: + """Tests for register_json with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.register_json("t", "s3://bucket/data.json", object_store=mock_store) + mock_register.assert_called_once_with("s3://", mock_store, host="bucket") + + +class TestReadJsonObjectStore: + """Tests for read_json with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.read_json("s3://bucket/data.json", object_store=mock_store) + mock_register.assert_called_once_with("s3://", mock_store, host="bucket") + + +class TestRegisterAvroObjectStore: + """Tests for register_avro with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.register_avro("t", "s3://bucket/data.avro", object_store=mock_store) + mock_register.assert_called_once_with("s3://", mock_store, host="bucket") + + +class TestReadAvroObjectStore: + """Tests for read_avro with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.read_avro("s3://bucket/data.avro", object_store=mock_store) + mock_register.assert_called_once_with("s3://", mock_store, host="bucket") + + +class TestRegisterArrowObjectStore: + """Tests for register_arrow with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.register_arrow( + "t", "s3://bucket/data.arrow", object_store=mock_store + ) + mock_register.assert_called_once_with("s3://", mock_store, host="bucket") + + +class TestReadArrowObjectStore: + """Tests for read_arrow with object_store parameter.""" + + def test_object_store_triggers_registration(self, ctx): + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + with contextlib.suppress(Exception): + ctx.read_arrow("s3://bucket/data.arrow", object_store=mock_store) + mock_register.assert_called_once_with("s3://", mock_store, host="bucket") + + +class TestEndToEndWithLocalFileSystem: + """Integration test using LocalFileSystem object store with register_parquet.""" + + def test_register_parquet_with_local_object_store(self, ctx, tmp_path): + """Verify the full flow works with a real object store and local file.""" + import pyarrow.parquet as pq + from datafusion.object_store import LocalFileSystem + + # Write a test parquet file + table = pa.table({"x": [1, 2, 3], "y": ["a", "b", "c"]}) + parquet_path = tmp_path / "test.parquet" + pq.write_table(table, str(parquet_path)) + + # Use file:// URL with LocalFileSystem object store + store = LocalFileSystem() + file_url = parquet_path.as_uri() + + ctx.register_parquet("test_tbl", file_url, object_store=store) + result = ctx.sql("SELECT * FROM test_tbl").collect() + + assert len(result) == 1 + assert result[0].num_rows == 3 + assert result[0].column("x").to_pylist() == [1, 2, 3] + + def test_read_parquet_with_local_object_store(self, ctx, tmp_path): + """Verify read_parquet works with object_store parameter.""" + import pyarrow.parquet as pq + from datafusion.object_store import LocalFileSystem + + table = pa.table({"val": [10, 20, 30]}) + parquet_path = tmp_path / "read_test.parquet" + pq.write_table(table, str(parquet_path)) + + store = LocalFileSystem() + file_url = parquet_path.as_uri() + + df = ctx.read_parquet(file_url, object_store=store) + result = df.collect() + + assert len(result) == 1 + assert result[0].column("val").to_pylist() == [10, 20, 30] Review Comment: ```suggestion @pytest.mark.parametrize( ("path", "scheme", "host"), [ ("s3://my-bucket/path/file.parquet", "s3://", "my-bucket"), ("gs://my-gcs-bucket/data.parquet", "gs://", "my-gcs-bucket"), ("az://my-container/data.parquet", "az://", "my-container"), ("https://example.com/data.parquet";, "https://";, "example.com"), ("file:///tmp/data.parquet", "file://", None), ], ) def test_register_object_store_for_url(ctx, path, scheme, host): store = MagicMock() with patch.object(ctx, "register_object_store") as register: ctx._register_object_store_for_path(path, store) register.assert_called_once_with(scheme, store, host=host) @pytest.mark.parametrize( ("path", "error"), [ ("/local/path/file.parquet", "Cannot determine object store URL"), ("relative/path.parquet", "Cannot determine object store URL"), (Path("/local/file.parquet"), "Cannot determine object store URL"), ("C:\\Users\\data\\file.parquet", "must include a host or bucket"), ("s3:///key.parquet", "must include a host or bucket"), ], ) def test_register_object_store_rejects_invalid_url(path, error, ctx): with pytest.raises(ValueError, match=error): ctx._register_object_store_for_path(path, MagicMock()) @pytest.mark.parametrize( ("method_name", "args", "path"), [ ("register_parquet", ("table",), "s3://bucket/data.parquet"), ("read_parquet", (), "s3://bucket/data.parquet"), ("register_csv", ("table",), "s3://bucket/data.csv"), ("read_csv", (), "s3://bucket/data.csv"), ("register_json", ("table",), "s3://bucket/data.json"), ("read_json", (), "s3://bucket/data.json"), ("register_avro", ("table",), "s3://bucket/data.avro"), ("read_avro", (), "s3://bucket/data.avro"), ("register_arrow", ("table",), "s3://bucket/data.arrow"), ("read_arrow", (), "s3://bucket/data.arrow"), ], ) def test_file_methods_register_object_store(ctx, method_name, args, path): store = MagicMock() # The remote file does not exist. Registration happens before DataFusion # tries to inspect it, which is the behavior under test. with ( patch.object(ctx, "register_object_store") as register, contextlib.suppress(Exception), ): getattr(ctx, method_name)(*args, path, object_store=store) register.assert_called_once_with("s3://", store, host="bucket") def test_register_csv_uses_first_path_for_object_store(ctx): store = MagicMock() paths = ["s3://bucket/a.csv", "s3://bucket/b.csv"] with ( patch.object(ctx, "register_object_store") as register, contextlib.suppress(Exception), ): ctx.register_csv("table", paths, object_store=store) register.assert_called_once_with("s3://", store, host="bucket") def test_object_store_none_does_not_register(ctx): with ( patch.object(ctx, "register_object_store") as register, contextlib.suppress(Exception), ): ctx.register_parquet("table", "missing.parquet") register.assert_not_called() def test_file_method_rejects_local_path_with_object_store(ctx): with pytest.raises(ValueError, match="Cannot determine object store URL"): ctx.register_parquet("table", "/local/file.parquet", object_store=MagicMock()) @pytest.mark.parametrize("method_name", ["register_parquet", "read_parquet"]) def test_parquet_methods_with_local_object_store(ctx, tmp_path, method_name): table = pa.table({"value": [10, 20, 30]}) parquet_path = tmp_path / "data.parquet" pq.write_table(table, parquet_path) path = parquet_path.as_uri() if method_name == "register_parquet": ctx.register_parquet("test_table", path, object_store=LocalFileSystem()) dataframe = ctx.sql("SELECT * FROM test_table") else: dataframe = ctx.read_parquet(path, object_store=LocalFileSystem()) assert dataframe.collect()[0].column("value").to_pylist() == [10, 20, 30] ``` I think we can use parameterization for the unit tests to simplify them. I think the above suggestion is more in line with the rest of the repository. I don't see why we'd need or want classes for the tests. I'm open to a counter argument, though. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
