airfRANS-model-exploration/tests/test_public_data.py

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