airfRANS-model-exploration/tests/test_public_data.py

from __future__ import annotations

import hashlib
import json
import gzip
import http.client
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.bounded_public import prepare_public_airfrans_processed_hf_bounded
from airfrans_frontier.raw.public import (
    HttpRangeReader,
    _extract_remote_case_members,
    _read_zip_central_directory,
    _remote_archive_case_members,
    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_range_streaming_extracts_contiguous_case_with_one_payload_read(self) -> None:
        class CountingRangeReader:
            def __init__(self, path: Path) -> None:
                self._path = path
                self.size = path.stat().st_size
                self.bytes_read = 0
                self.calls: list[tuple[int, int]] = []

            def read_range(self, start: int, length: int) -> bytes:
                self.calls.append((start, length))
                with self._path.open("rb") as handle:
                    handle.seek(start)
                    data = handle.read(length)
                self.bytes_read += len(data)
                return data

        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])
            reader = CountingRangeReader(archive)
            members = _read_zip_central_directory(reader)
            case_members = _remote_archive_case_members(members)[case_name]

            reader.calls.clear()
            _extract_remote_case_members(reader, case_members, tmp_path / "streaming_raw")

            self.assertEqual(len(reader.calls), 1)
            self.assertTrue((tmp_path / "streaming_raw" / case_name / "constant" / "transportProperties").is_file())

    def test_http_range_reader_retries_timeout_before_failing_run(self) -> None:
        class FakeResponse:
            status = 206

            def __enter__(self):
                return self

            def __exit__(self, exc_type, exc, traceback):
                return False

            def read(self) -> bytes:
                return b"ok"

        reader = HttpRangeReader.__new__(HttpRangeReader)
        reader.url = "https://example.test/OF_dataset.zip"
        reader.size = 10
        reader.bytes_read = 0

        with patch("airfrans_frontier.raw.public.urllib.request.urlopen", side_effect=[TimeoutError("timed out"), FakeResponse()]), patch(
            "airfrans_frontier.raw.public.time.sleep"
        ) as sleep:
            data = reader.read_range(2, 2)

        self.assertEqual(data, b"ok")
        self.assertEqual(reader.bytes_read, 2)
        sleep.assert_called_once_with(2.0)

    def test_http_range_reader_retries_incomplete_body_before_failing_run(self) -> None:
        class FakeResponse:
            status = 206

            def __enter__(self):
                return self

            def __exit__(self, exc_type, exc, traceback):
                return False

            def read(self) -> bytes:
                return b"ok"

        reader = HttpRangeReader.__new__(HttpRangeReader)
        reader.url = "https://example.test/OF_dataset.zip"
        reader.size = 10
        reader.bytes_read = 0

        with patch(
            "airfrans_frontier.raw.public.urllib.request.urlopen",
            side_effect=[http.client.IncompleteRead(b"pa", 2), FakeResponse()],
        ), patch("airfrans_frontier.raw.public.time.sleep") as sleep:
            data = reader.read_range(2, 2)

        self.assertEqual(data, b"ok")
        self.assertEqual(reader.bytes_read, 2)
        sleep.assert_called_once_with(2.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)


    def test_bounded_public_hf_uploads_verified_chunks_and_cleans_staging(self) -> None:
        uploaded: dict[str, tuple[int, str, bytes]] = {}
        commits: list[tuple[str, tuple[str, ...]]] = []

        class FakeCommitOperationAdd:
            def __init__(self, *, path_in_repo: str, path_or_fileobj: str) -> None:
                self.path_in_repo = path_in_repo
                self.path_or_fileobj = path_or_fileobj

        class FakeApi:
            def __init__(self, token: str) -> None:
                self.token = token

            def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None:
                self.repo_id = repo_id

            def create_commit(self, *, repo_id: str, repo_type: str, operations, commit_message: str):
                paths: list[str] = []
                for operation in operations:
                    payload = Path(operation.path_or_fileobj).read_bytes()
                    uploaded[operation.path_in_repo] = (
                        len(payload),
                        hashlib.sha256(payload).hexdigest(),
                        payload,
                    )
                    paths.append(operation.path_in_repo)
                commits.append((commit_message, tuple(paths)))
                return types.SimpleNamespace(commit_url=f"https://huggingface.co/datasets/{repo_id}/commit/{len(commits)}", oid=str(len(commits)))

            def repo_info(self, *, repo_id: str, repo_type: str, files_metadata: bool):
                siblings = [
                    types.SimpleNamespace(rfilename=path, size=size, lfs={"sha256": sha})
                    for path, (size, sha, _payload) in uploaded.items()
                ]
                return types.SimpleNamespace(siblings=siblings)

            def list_repo_files(self, *, repo_id: str, repo_type: str):
                return sorted(uploaded)

        fake_module = types.SimpleNamespace(HfApi=FakeApi, CommitOperationAdd=FakeCommitOperationAdd)
        with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(
            "os.environ", {"HF_TOKEN": "token"}
        ):
            tmp_path = Path(tmp)
            archive = tmp_path / "OF_dataset.zip"
            case_names = ["airFoil2D_SST_10.0_5.0_0012", "airFoil2D_SST_11.0_5.0_0012"]
            write_minimal_airfrans_archive(archive, case_names)

            report = prepare_public_airfrans_processed_hf_bounded(
                repo_id="owner/repo",
                path_in_repo="processed/full",
                work_dir=tmp_path / "work",
                output_dir=tmp_path / "staging",
                source_url=str(archive),
                min_cases=2,
                chunk_max_bytes=1,
                train_cases=1,
                val_cases=1,
                test_cases=0,
                split_seed=123,
            )

            state = json.loads((tmp_path / "work" / "bounded_prepare_state.json").read_text())
            final_manifest = json.loads((tmp_path / "work" / "chunk_manifests" / "hf_dataset_manifest.json").read_text())

            self.assertTrue(report["ok"])
            self.assertEqual(report["processed_case_count"], 2)
            self.assertEqual(report["chunks_uploaded"], 2)
            self.assertEqual(state["phase"], "published")
            self.assertEqual(final_manifest["case_count"], 2)
            self.assertEqual(final_manifest["feature_names"][0], "x")
            self.assertEqual(final_manifest["target_names"], ["velocity_x", "velocity_y", "pressure", "turbulent_viscosity"])
            self.assertEqual(final_manifest["split_compatibility"]["train_cases"], 1)
            self.assertFalse(any((tmp_path / "staging").glob("*.npz")))
            self.assertFalse((tmp_path / "work" / "bounded_raw_scratch").exists())
            self.assertIn("processed/full/hf_dataset_manifest.json", uploaded)
            self.assertEqual(sum(1 for path in uploaded if path.endswith(".npz")), 2)
            self.assertEqual(sum(1 for message, _paths in commits if message.startswith("Upload bounded AirfRANS processed chunk")), 2)

    def test_bounded_public_hf_rejects_remote_checksum_mismatch(self) -> None:
        uploaded: dict[str, tuple[int, str]] = {}

        class FakeCommitOperationAdd:
            def __init__(self, *, path_in_repo: str, path_or_fileobj: str) -> None:
                self.path_in_repo = path_in_repo
                self.path_or_fileobj = path_or_fileobj

        class FakeApi:
            def __init__(self, token: str) -> None:
                self.token = token

            def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None:
                pass

            def create_commit(self, *, repo_id: str, repo_type: str, operations, commit_message: str):
                for operation in operations:
                    payload = Path(operation.path_or_fileobj).read_bytes()
                    uploaded[operation.path_in_repo] = (len(payload), hashlib.sha256(payload).hexdigest())
                return types.SimpleNamespace(commit_url="https://huggingface.co/datasets/owner/repo/commit/bad", oid="bad")

            def repo_info(self, *, repo_id: str, repo_type: str, files_metadata: bool):
                siblings = [
                    types.SimpleNamespace(rfilename=path, size=size, lfs={"sha256": "0" * 64})
                    for path, (size, _sha) in uploaded.items()
                ]
                return types.SimpleNamespace(siblings=siblings)

        fake_module = types.SimpleNamespace(HfApi=FakeApi, CommitOperationAdd=FakeCommitOperationAdd)
        with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(
            "os.environ", {"HF_TOKEN": "token"}
        ):
            tmp_path = Path(tmp)
            archive = tmp_path / "OF_dataset.zip"
            write_minimal_airfrans_archive(archive, ["airFoil2D_SST_10.0_5.0_0012"])

            with self.assertRaisesRegex(RuntimeError, "HF upload verification failed"):
                prepare_public_airfrans_processed_hf_bounded(
                    repo_id="owner/repo",
                    path_in_repo="processed/full",
                    work_dir=tmp_path / "work",
                    output_dir=tmp_path / "staging",
                    source_url=str(archive),
                    min_cases=1,
                    chunk_max_bytes=1024,
                    train_cases=1,
                    val_cases=0,
                    test_cases=0,
                    split_seed=123,
                )

if __name__ == "__main__":
    unittest.main()