68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
import types
|
||
|
|
import unittest
|
||
|
|
import zipfile
|
||
|
|
from pathlib import Path
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
from airfrans_frontier.raw.public import ensure_public_airfrans_processed_hf, extract_of_dataset
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|