backend/tests/test_bootstrap.py

"""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"