diff --git a/tests/core/test_batch.py b/tests/core/test_batch.py index 740f6d281e..bb586ef45a 100644 --- a/tests/core/test_batch.py +++ b/tests/core/test_batch.py @@ -32,9 +32,33 @@ def test_storage_get_url(): assert url.replace("\\", "/").endswith("root_dir/prefix--file.jsonl.gz") -def test_storage_from_url(): - url = urlparse("s3://bucket/path/to/file?region=us-east-1") +@pytest.mark.parametrize( + "file_url,root,prefix,params", + [ + pytest.param( + "s3://bucket/path/to/file?region=us-east-1", + "s3://bucket", + None, + {"region": ["us-east-1"]}, + id="s3", + ), + pytest.param( + "file://abs/path/to/file", + "file://abs", + None, + {}, + id="local", + ), + ], +) +def test_storage_from_url(file_url: str, root: str, prefix: str, params: dict): + """Test storage target from URL.""" + url = urlparse(file_url) target = StorageTarget.from_url(url) - assert target.root == "s3://bucket" - assert target.prefix is None - assert target.params == {"region": ["us-east-1"]} + assert target.root == root + assert target.prefix == prefix + assert target.params == params + + if url.scheme == "file": + with target.fs(create=True) as fs: + assert fs.geturl("file.jsonl.gz").startswith(root)