bc76475087
Co-authored-by: aider (openrouter/qwen/qwen3-coder) <aider@aider.chat>
383 lines
13 KiB
Python
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
|