Files
dataset-uploader/features/steps/dataset_validation_steps.py
brent.edwards bc76475087 test: add tests for Bloom filter subset functionality
Co-authored-by: aider (openrouter/qwen/qwen3-coder) <aider@aider.chat>
2026-01-28 18:52:21 -08:00

383 lines
13 KiB
Python

"""Step definitions for dataset validation tests."""
from __future__ import annotations
import csv
import shutil
import sys
from pathlib import Path
from types import SimpleNamespace
from typing import TYPE_CHECKING
from unittest.mock import MagicMock
import pandas as pd
from behave import given, then, when
from bloom_filter2 import BloomFilter
SCRIPTS_DIR = Path(__file__).parent.parent.parent / "scripts"
sys.path.insert(0, str(SCRIPTS_DIR))
from dataset_validator import ( # noqa: E402
collect_parquet_strings,
collect_source_strings,
dataset_can_download_from_hf,
dataset_exists_on_hf,
dataset_has_required_columns,
dataset_is_parquet_file,
)
if TYPE_CHECKING:
from behave.runner import Context
def _ensure_parquet_file(context: Context) -> tuple[Path, bool]:
file_path = getattr(context, "dataset_file_path", None)
if file_path is not None:
return file_path, False
base_dir = getattr(context, "scenario_temp_dir", Path("/tmp"))
base_dir.mkdir(parents=True, exist_ok=True)
file_path = base_dir / "example.parquet"
df = pd.DataFrame(
[
{
"subject": "http://example.org/s",
"predicate": "http://example.org/p",
"object": "http://example.org/o",
"object_type": "uri",
"object_datatype": None,
"object_language": None,
}
]
)
df.to_parquet(file_path, index=False)
context.dataset_file_path = file_path
return file_path, True
def _ensure_validation_base_dir(context: Context) -> Path:
base_dir = getattr(context, "validation_base_dir", None)
if base_dir is not None:
return base_dir
scenario_base = getattr(context, "scenario_temp_dir", Path("/tmp"))
base_dir = scenario_base / "dataset_validation"
base_dir.mkdir(parents=True, exist_ok=True)
context.validation_base_dir = base_dir
return base_dir
@given('a dataset id "{dataset_id}"')
def step_set_dataset_id(context: Context, dataset_id: str) -> None:
"""Set the dataset id for validation."""
context.dataset_id = dataset_id
@given('a dataset file path "{file_path}"')
def step_set_dataset_file_path(context: Context, file_path: str) -> None:
"""Set the dataset file path for validation."""
context.dataset_file_path = Path(file_path)
@given("the HuggingFace API reports the dataset exists")
def step_set_hf_api_exists(context: Context) -> None:
"""Configure the mock to report that the dataset exists."""
context.hf_api_exists = True
@given("the HuggingFace download succeeds")
def step_set_hf_download_succeeds(context: Context) -> None:
"""Configure the mock to report that the download succeeds."""
context.hf_download_succeeds = True
@given("the dataset file can be read as parquet")
def step_set_parquet_readable(context: Context) -> None:
"""Configure the mock to report that the file is readable as parquet."""
context.parquet_read_succeeds = True
@given("a dataset download directory with CSV and N-Triples data")
def step_create_download_dir_with_data(context: Context) -> None:
"""Create a dataset download directory with sample CSV and N-Triples files."""
base_dir = _ensure_validation_base_dir(context)
dataset_id = getattr(context, "dataset_id", "example-dataset")
download_dir = base_dir / "downloads" / dataset_id
download_dir.mkdir(parents=True, exist_ok=True)
csv_path = download_dir / "sample.csv"
with open(csv_path, "w", newline="", encoding="utf-8") as file_obj:
writer = csv.writer(file_obj)
writer.writerow(["hello", "world"])
writer.writerow([" ", ""])
nt_path = download_dir / "sample.nt"
nt_path.write_text(
'<http://example.org/s> <http://example.org/p> "object" .\n',
encoding="utf-8",
)
context.cleanup_download_dir = True
context.download_dir = download_dir
@given("a parquet file with sample data")
def step_create_parquet_with_sample_data(context: Context) -> None:
"""Create a parquet file with sample RDF-style columns."""
base_dir = getattr(context, "scenario_temp_dir", Path("/tmp"))
base_dir.mkdir(parents=True, exist_ok=True)
file_path = base_dir / "sample.parquet"
df = pd.DataFrame(
[
{
"subject": "http://example.org/s",
"predicate": "http://example.org/p",
"object": "object",
"object_type": "literal",
"object_datatype": None,
"object_language": " ",
}
]
)
df.to_parquet(file_path, index=False)
context.dataset_file_path = file_path
context.parquet_file_created = True
@given("two bloom filters with the same element")
def step_create_bloom_filters_same_element(context: Context) -> None:
"""Create two bloom filters and add the same element to both."""
context.bf1 = BloomFilter(max_elements=1000, error_rate=0.1)
context.bf2 = BloomFilter(max_elements=1000, error_rate=0.1)
element = "test_element"
context.bf1.add(element)
context.bf2.add(element)
context.element = element
@given("two bloom filters with one shared element and one unique to the second")
def step_create_bloom_filters_partial_overlap(context: Context) -> None:
"""Create two bloom filters with shared and unique elements."""
context.bf1 = BloomFilter(max_elements=1000, error_rate=0.1)
context.bf2 = BloomFilter(max_elements=1000, error_rate=0.1)
shared_element = "shared_element"
unique_element = "unique_element"
context.bf1.add(shared_element)
context.bf2.add(shared_element)
context.bf2.add(unique_element)
context.shared_element = shared_element
context.unique_element = unique_element
@given("two bloom filters with different elements")
def step_create_bloom_filters_different_elements(context: Context) -> None:
"""Create two bloom filters with completely different elements."""
context.bf1 = BloomFilter(max_elements=1000, error_rate=0.1)
context.bf2 = BloomFilter(max_elements=1000, error_rate=0.1)
element1 = "element_1"
element2 = "element_2"
context.bf1.add(element1)
context.bf2.add(element2)
context.element1 = element1
context.element2 = element2
@when("I validate dataset existence")
def step_validate_dataset_existence(context: Context) -> None:
"""Validate dataset existence using the validator."""
mock_api = MagicMock()
if getattr(context, "hf_api_exists", False):
mock_api.dataset_info.return_value = SimpleNamespace(id=context.dataset_id)
else:
from huggingface_hub.errors import HfHubHTTPError
mock_api.dataset_info.side_effect = HfHubHTTPError(
"Not Found",
response=SimpleNamespace(status_code=404),
)
context.dataset_exists = dataset_exists_on_hf(
context.dataset_id,
organization="test-org",
api=mock_api,
)
@when("I validate dataset download")
def step_validate_dataset_download(context: Context) -> None:
"""Validate dataset download using the validator."""
download_func = MagicMock()
if getattr(context, "hf_download_succeeds", False):
download_func.return_value = "/tmp/hf-download"
else:
from huggingface_hub.errors import HfHubHTTPError
download_func.side_effect = HfHubHTTPError(
"Not Found",
response=SimpleNamespace(status_code=404),
)
context.dataset_downloadable = bool(
dataset_can_download_from_hf(
context.dataset_id,
organization="test-org",
download_func=download_func,
)
)
@when("I validate the dataset file format")
def step_validate_dataset_file_format(context: Context) -> None:
"""Validate dataset file format using the validator."""
file_path, created_file = _ensure_parquet_file(context)
try:
context.dataset_is_parquet = dataset_is_parquet_file(file_path)
finally:
if created_file and file_path.exists():
file_path.unlink()
@when("I validate the dataset columns")
def step_validate_dataset_columns(context: Context) -> None:
"""Validate dataset file columns using the validator."""
file_path, created_file = _ensure_parquet_file(context)
try:
context.dataset_has_columns = dataset_has_required_columns(file_path)
finally:
if created_file and file_path.exists():
file_path.unlink()
@when("I collect source strings")
def step_collect_source_strings(context: Context) -> None:
"""Collect strings from the source dataset files."""
base_dir = _ensure_validation_base_dir(context)
dataset_id = getattr(context, "dataset_id", "example-dataset")
context.source_strings = collect_source_strings(dataset_id, base_dir)
if getattr(context, "cleanup_download_dir", False):
download_dir = getattr(context, "download_dir", None)
if download_dir and download_dir.exists():
shutil.rmtree(download_dir)
@when("I collect parquet strings")
def step_collect_parquet_strings(context: Context) -> None:
"""Collect strings from a parquet file."""
file_path, created_file = _ensure_parquet_file(context)
try:
context.parquet_strings = collect_parquet_strings(file_path)
finally:
if created_file and file_path.exists():
file_path.unlink()
if getattr(context, "parquet_file_created", False) and file_path.exists():
file_path.unlink()
@when("I check if the first bloom filter is a subset of the second")
def step_check_bloom_filter_subset(context: Context) -> None:
"""Check if the first bloom filter is a subset of the second."""
# Import the function from upload_all_datasets.py
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
from upload_all_datasets import is_sub_bloom_filter
context.subset_result = is_sub_bloom_filter(context.bf1, context.bf2)
@when("I check if the second bloom filter is a subset of the first")
def step_check_bloom_filter_subset_reverse(context: Context) -> None:
"""Check if the second bloom filter is a subset of the first."""
# Import the function from upload_all_datasets.py
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
from upload_all_datasets import is_sub_bloom_filter
context.subset_result_reverse = is_sub_bloom_filter(context.bf2, context.bf1)
@then("the dataset should be reported as existing")
def step_assert_dataset_exists(context: Context) -> None:
"""Assert that the dataset exists."""
assert context.dataset_exists is True
@then("the dataset should be reported as downloadable")
def step_assert_dataset_downloadable(context: Context) -> None:
"""Assert that the dataset is downloadable."""
assert context.dataset_downloadable is True
@then("the dataset file should be reported as parquet")
def step_assert_dataset_is_parquet(context: Context) -> None:
"""Assert that the dataset file is parquet."""
assert context.dataset_is_parquet is True
@then("the dataset should have the required columns")
def step_assert_dataset_has_columns(context: Context) -> None:
"""Assert that the dataset has required columns."""
assert context.dataset_has_columns is True
@then('the source strings should include "{value}"')
def step_assert_source_strings_include(context: Context, value: str) -> None:
"""Assert that a value is in the source string set."""
assert value in context.source_strings
@then('the source strings should not include "{value}"')
def step_assert_source_strings_exclude(context: Context, value: str) -> None:
"""Assert that a value is not in the source string set."""
assert value not in context.source_strings
@then('the parquet strings should include "{value}"')
def step_assert_parquet_strings_include(context: Context, value: str) -> None:
"""Assert that a value is in the parquet string set."""
assert value in context.parquet_strings
@then('the parquet strings should not include "{value}"')
def step_assert_parquet_strings_exclude(context: Context, value: str) -> None:
"""Assert that a value is not in the parquet string set."""
assert value not in context.parquet_strings
@then("the first bloom filter should be a subset of the second")
def step_assert_first_is_subset_of_second(context: Context) -> None:
"""Assert that the first bloom filter is a subset of the second."""
assert context.subset_result is True
@then("the first bloom filter should not be a subset of the second")
def step_assert_first_is_not_subset_of_second(context: Context) -> None:
"""Assert that the first bloom filter is not a subset of the second."""
assert context.subset_result is False
@then("the second bloom filter should be a subset of the first")
def step_assert_second_is_subset_of_first(context: Context) -> None:
"""Assert that the second bloom filter is a subset of the first."""
assert context.subset_result_reverse is True
@then("the second bloom filter should not be a subset of the first")
def step_assert_second_is_not_subset_of_first(context: Context) -> None:
"""Assert that the second bloom filter is not a subset of the first."""
assert context.subset_result_reverse is False