161 lines
7.4 KiB
Python
161 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
import gzip
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import shutil
|
|
import unittest
|
|
import zipfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from airfrans_frontier.runtime import remove_pythonpath_entries
|
|
|
|
remove_pythonpath_entries()
|
|
|
|
from airfrans_frontier.raw.public import ensure_public_airfrans_processed_hf, extract_of_dataset, process_of_dataset_url_streaming
|
|
|
|
|
|
def write_minimal_airfrans_archive(archive: Path, case_names: list[str]) -> None:
|
|
with zipfile.ZipFile(archive, "w") as zf:
|
|
for case_name in case_names:
|
|
base = f"OF_dataset/{case_name}"
|
|
zf.writestr(f"{base}/constant/transportProperties", "nu 1e-5;\n")
|
|
zf.writestr(
|
|
f"{base}/constant/polyMesh/boundary",
|
|
"\naerofoil\n{\n type wall;\n nFaces 1;\n startFace 0;\n}\nfarfield\n{\n type patch;\n nFaces 3;\n startFace 1;\n}\n",
|
|
)
|
|
zf.writestr(f"{base}/constant/polyMesh/points.gz", gzip.compress(b"4\n(\n(0 0 0)\n(1 0 0)\n(1 1 0)\n(0 1 0)\n)\n"))
|
|
zf.writestr(f"{base}/constant/polyMesh/faces.gz", gzip.compress(b"4\n(\n2(0 1)\n2(1 2)\n2(2 3)\n2(3 0)\n)\n"))
|
|
zf.writestr(f"{base}/constant/polyMesh/owner.gz", gzip.compress(b"4\n(\n0\n0\n0\n0\n)\n"))
|
|
zf.writestr(f"{base}/constant/polyMesh/neighbour.gz", gzip.compress(b"0\n(\n)\n"))
|
|
zf.writestr(f"{base}/1/U.gz", gzip.compress(b"1\n(\n(1 0 0)\n)\n"))
|
|
zf.writestr(f"{base}/1/p.gz", gzip.compress(b"1\n(\n0.5\n)\n"))
|
|
zf.writestr(f"{base}/1/nut.gz", gzip.compress(b"1\n(\n0.01\n)\n"))
|
|
|
|
|
|
class PublicAirfransDataTests(unittest.TestCase):
|
|
def test_prepare_public_hf_skips_when_dataset_already_published(self) -> None:
|
|
class FakeApi:
|
|
def __init__(self, token=None):
|
|
self.token = token
|
|
|
|
def list_repo_files(self, *, repo_id: str, repo_type: str):
|
|
assert repo_id == "owner/airfrans-processed"
|
|
assert repo_type == "dataset"
|
|
return [
|
|
"processed/full/case_000.npz",
|
|
"processed/full/case_001.npz",
|
|
"processed/full/hf_dataset_manifest.json",
|
|
]
|
|
|
|
fake_module = types.SimpleNamespace(HfApi=FakeApi)
|
|
with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(
|
|
"os.environ", {"HF_TOKEN": "token"}
|
|
):
|
|
report = ensure_public_airfrans_processed_hf(
|
|
repo_id="owner/airfrans-processed",
|
|
path_in_repo="processed/full",
|
|
work_dir=Path(tmp) / "work",
|
|
output_dir=Path(tmp) / "out",
|
|
min_cases=2,
|
|
)
|
|
|
|
self.assertTrue(report["ok"])
|
|
self.assertEqual(report["phase"], "already_published")
|
|
self.assertEqual(report["npz_file_count"], 2)
|
|
self.assertTrue(report["has_manifest"])
|
|
|
|
def test_extract_of_dataset_finds_public_archive_root(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
archive = tmp_path / "OF_dataset.zip"
|
|
with zipfile.ZipFile(archive, "w") as zf:
|
|
zf.writestr("OF_dataset/airFoil2D_SST_demo/system/controlDict", "ok")
|
|
root = extract_of_dataset(archive, tmp_path / "raw", min_cases=1)
|
|
|
|
self.assertEqual(root.name, "OF_dataset")
|
|
|
|
def test_extract_of_dataset_rejects_zip_slip_paths(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
archive = tmp_path / "bad.zip"
|
|
with zipfile.ZipFile(archive, "w") as zf:
|
|
zf.writestr("../escape.txt", "bad")
|
|
with self.assertRaisesRegex(RuntimeError, "Unsafe path"):
|
|
extract_of_dataset(archive, tmp_path / "raw", min_cases=1)
|
|
|
|
def test_extract_of_dataset_fails_before_partial_extract_when_disk_is_too_small(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
archive = tmp_path / "OF_dataset.zip"
|
|
with zipfile.ZipFile(archive, "w") as zf:
|
|
zf.writestr("OF_dataset/airFoil2D_SST_demo/system/controlDict", "ok")
|
|
tiny_disk = shutil._ntuple_diskusage(total=10, used=10, free=0)
|
|
with patch("airfrans_frontier.raw.public.shutil.disk_usage", return_value=tiny_disk):
|
|
with self.assertRaisesRegex(RuntimeError, "Insufficient free disk"):
|
|
extract_of_dataset(archive, tmp_path / "raw", min_cases=1)
|
|
self.assertFalse((tmp_path / "raw" / "OF_dataset").exists())
|
|
|
|
def test_range_streaming_processing_writes_npz_and_discards_raw_case(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
archive = tmp_path / "OF_dataset.zip"
|
|
case_name = "airFoil2D_SST_10.0_5.0_0012"
|
|
write_minimal_airfrans_archive(archive, [case_name])
|
|
|
|
streamed = process_of_dataset_url_streaming(
|
|
str(archive),
|
|
tmp_path / "processed",
|
|
scratch_dir=tmp_path / "streaming_raw",
|
|
min_cases=1,
|
|
progress_every=1,
|
|
)
|
|
result = streamed.processing
|
|
|
|
self.assertEqual(result.case_count, 1)
|
|
self.assertTrue((tmp_path / "processed" / f"{case_name}.npz").is_file())
|
|
self.assertTrue(result.manifest_path.is_file())
|
|
self.assertFalse((tmp_path / "streaming_raw" / case_name).exists())
|
|
self.assertGreater(streamed.ranged_bytes_read, 0)
|
|
|
|
def test_prepare_public_hf_streams_archive_before_publish(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
source_archive = tmp_path / "source_OF_dataset.zip"
|
|
case_name = "airFoil2D_SST_10.0_5.0_0012"
|
|
write_minimal_airfrans_archive(source_archive, [case_name])
|
|
statuses = [
|
|
{"file_count": 0, "npz_file_count": 0, "has_manifest": False},
|
|
{"file_count": 2, "npz_file_count": 1, "has_manifest": True},
|
|
]
|
|
|
|
def fake_publish(**kwargs):
|
|
data_root = Path(kwargs["data_root"])
|
|
self.assertTrue((data_root / f"{case_name}.npz").is_file())
|
|
self.assertFalse((tmp_path / "work" / "streaming_raw" / case_name).exists())
|
|
return {"repo_url": "https://huggingface.co/datasets/owner/repo", "npz_file_count": 1}
|
|
|
|
with patch("airfrans_frontier.raw.public._hf_dataset_status", side_effect=statuses), patch(
|
|
"airfrans_frontier.raw.public.publish_processed_dataset", side_effect=fake_publish
|
|
):
|
|
report = ensure_public_airfrans_processed_hf(
|
|
repo_id="owner/repo",
|
|
path_in_repo="processed/full",
|
|
work_dir=tmp_path / "work",
|
|
output_dir=tmp_path / "processed",
|
|
source_url=str(source_archive),
|
|
min_cases=1,
|
|
)
|
|
self.assertFalse((tmp_path / "work" / "OF_dataset.zip").exists())
|
|
|
|
self.assertTrue(report["ok"])
|
|
self.assertTrue(report["streaming"])
|
|
self.assertEqual(report["streaming_mode"], "zip_range")
|
|
self.assertEqual(report["download"]["mode"], "zip_range")
|
|
self.assertEqual(report["processed_case_count"], 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|