Skip to content

Commit

Permalink
Improve LocationProvider unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Sreesh Maheshwar committed Jan 12, 2025
1 parent cad0ad7 commit 216e9a9
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions tests/table/test_locations.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,18 @@ def new_data_location(self, data_file_name: str, partition_key: Optional[Partiti
return f"custom_location_provider/{data_file_name}"


def test_default_location_provider() -> None:
def test_simple_location_provider_no_partition() -> None:
provider = load_location_provider(table_location="table_location", table_properties=EMPTY_DICT)

assert provider.new_data_location("my_file") == "table_location/data/my_file"


def test_simple_location_provider_with_partition() -> None:
provider = load_location_provider(table_location="table_location", table_properties=EMPTY_DICT)

assert provider.new_data_location("my_file", PARTITION_KEY) == "table_location/data/string_field=example_string/my_file"


def test_custom_location_provider() -> None:
qualified_name = CustomLocationProvider.__module__ + "." + CustomLocationProvider.__name__
provider = load_location_provider(
Expand Down Expand Up @@ -82,19 +88,17 @@ def test_object_storage_injects_entropy() -> None:
assert all(c in "01" for c in dir_name)


@pytest.mark.parametrize("object_storage", [True, False])
def test_partition_value_in_path(object_storage: bool) -> None:
def test_object_storage_with_partition() -> None:
provider = load_location_provider(
table_location="table_location",
table_properties={
"write.object-storage.enabled": str(object_storage),
},
table_properties={"write.object-storage.enabled": "true"},
)

location = provider.new_data_location("test.parquet", PARTITION_KEY)
partition_segment = location.split("/")[-2]

assert partition_segment == "string_field=example_string"
# Partition values AND entropy included in the path. Hash differs to in the test below because partition key AND
# data file name are used as the hash input.
assert location == "table_location/data/0001/0010/1001/00000011/string_field=example_string/test.parquet"


# NB: We test here with None partition key too because disabling partitioned paths still replaces final / with - even in
Expand Down

0 comments on commit 216e9a9

Please sign in to comment.