| """Tests for astroparse_api.bootstrap.ensure_corpus.""" |
| from pathlib import Path |
|
|
| import astroparse_api.bootstrap as bmod |
| from astroparse_api.bootstrap import ensure_corpus |
|
|
|
|
| |
| |
| |
|
|
| def test_noop_when_both_exist(tmp_path, monkeypatch): |
| """ensure_corpus must not call snapshot_download when corpus is present.""" |
| |
| 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" |
|
|
|
|
| |
| |
| |
|
|
| 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" |
| |
|
|
| 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" |
|
|
