"""Tests for astroparse_api.bootstrap.ensure_corpus.""" from pathlib import Path import astroparse_api.bootstrap as bmod from astroparse_api.bootstrap import ensure_corpus # --------------------------------------------------------------------------- # No-op path: both artefacts exist → no download attempted # --------------------------------------------------------------------------- def test_noop_when_both_exist(tmp_path, monkeypatch): """ensure_corpus must not call snapshot_download when corpus is present.""" # Create the artefacts so the early-return guard fires corpus_dir = tmp_path / "corpus" corpus_dir.mkdir() faiss_path = tmp_path / "index.faiss" faiss_path.write_bytes(b"fake-faiss") called = [] def fake_snapshot(*args, **kwargs): called.append(True) raise AssertionError("snapshot_download should NOT be called when files exist") monkeypatch.setattr(bmod, "snapshot_download", fake_snapshot) ensure_corpus(corpus_dir, faiss_path, "kiyer/beacon_corpus") assert called == [], "snapshot_download was called despite artefacts existing" # --------------------------------------------------------------------------- # Download path: files missing → hub called, artefacts placed correctly # --------------------------------------------------------------------------- def test_download_when_both_missing(tmp_path, monkeypatch): """When both artefacts are absent, snapshot_download is called and files placed.""" corpus_dir = tmp_path / "corpus" faiss_path = tmp_path / "astroparse_fp16.faiss" assert not corpus_dir.exists() assert not faiss_path.exists() def fake_snapshot_download(repo_id, repo_type, local_dir, **kwargs): local_dir = Path(local_dir) (local_dir / "astroparse_fp16.faiss").write_bytes(b"faiss-data") corpus_sub = local_dir / "corpus" corpus_sub.mkdir(parents=True) (corpus_sub / "data-00000.parquet").write_bytes(b"parquet-data") monkeypatch.setattr(bmod, "snapshot_download", fake_snapshot_download) ensure_corpus(corpus_dir, faiss_path, "kiyer/beacon_corpus") assert faiss_path.is_file(), "FAISS index was not placed at faiss_path" assert faiss_path.read_bytes() == b"faiss-data" assert corpus_dir.is_dir(), "corpus_dir was not created" assert (corpus_dir / "data-00000.parquet").is_file(), "parquet shard missing" def test_download_when_faiss_missing(tmp_path, monkeypatch): """Download is triggered even when only the FAISS file is absent.""" corpus_dir = tmp_path / "corpus" corpus_dir.mkdir() (corpus_dir / "data-00000.parquet").write_bytes(b"existing-parquet") faiss_path = tmp_path / "astroparse_fp16.faiss" # faiss_path intentionally absent download_called = [] def fake_snapshot_download(repo_id, repo_type, local_dir, **kwargs): local_dir = Path(local_dir) (local_dir / "astroparse_fp16.faiss").write_bytes(b"new-faiss") sub = local_dir / "corpus" sub.mkdir() (sub / "data-00000.parquet").write_bytes(b"new-parquet") download_called.append(True) monkeypatch.setattr(bmod, "snapshot_download", fake_snapshot_download) ensure_corpus(corpus_dir, faiss_path, "kiyer/beacon_corpus") assert download_called, "expected snapshot_download to be called" assert faiss_path.is_file() assert faiss_path.read_bytes() == b"new-faiss"