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()