tests/test_hf_bootstrap.py

"""Persistence and optional Hub bootstrap behavior."""
from __future__ import annotations

import tarfile
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from backend import Backend


def test_user_data_persists_across_backend_instances_same_file(tmp_path: Path):
    """Same SQLite path = same rows after process-style reconnect (simulates redeploy with volume)."""
    db = tmp_path / "store.db"
    b1 = Backend(db)
    u = b1.create_user()
    uid = u.id
    del b1
    b2 = Backend(db)
    again = b2.get_user(uid)
    assert again is not None
    assert again["id"] == uid


def test_ensure_store_db_uses_existing_file(tmp_path: Path):
    from backend.hf_bootstrap import ensure_store_db

    p = tmp_path / "x.db"
    p.write_bytes(b"not-real-sqlite")
    assert ensure_store_db(p) == p.resolve()


def test_ensure_media_cache_downloads_and_extracts_archive(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
    from backend import hf_bootstrap

    remote_dir = tmp_path / "remote"
    source_dir = tmp_path / "source"
    source_dir.mkdir()
    (source_dir / "fetlife_fetishes").mkdir()
    (source_dir / "fetlife_fetishes" / "1.jpg").write_bytes(b"jpeg")
    archive = remote_dir / "cached_assets.tar"
    remote_dir.mkdir()
    with tarfile.open(archive, "w") as tar:
        tar.add(source_dir / "fetlife_fetishes" / "1.jpg", arcname="fetlife_fetishes/1.jpg")

    seen: dict[str, str] = {}

    def fake_download(*, repo_id, filename, revision, token, local_dir, repo_type, **kwargs):
        seen["repo_id"] = repo_id
        seen["filename"] = filename
        dest = Path(local_dir) / filename
        dest.write_bytes(archive.read_bytes())
        return str(dest)

    target = tmp_path / "cached_assets"
    monkeypatch.setenv("KINK_HF_MEDIA_ARCHIVE", "cached_assets.tar")
    monkeypatch.setenv("KINK_HF_MEDIA_DATASET_REPO", "dummy/media")
    with patch("huggingface_hub.hf_hub_download", side_effect=fake_download):
        out = hf_bootstrap.ensure_media_cache(target)

    assert out == target.resolve()
    assert seen == {"repo_id": "dummy/media", "filename": "cached_assets.tar"}
    assert (target / "fetlife_fetishes" / "1.jpg").read_bytes() == b"jpeg"
    assert (target / ".cached_assets.tar.ready").is_file()


def test_ensure_store_db_require_full_replaces_tiny_existing_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
    """Stale bundled seed on disk must not block Hub download when full catalog is required."""
    from backend import hf_bootstrap

    target = tmp_path / "kink_store.db"
    target.write_bytes(b"tiny")

    def fake_download(*, repo_id, filename, revision, token, local_dir, repo_type, **kwargs):
        dest = Path(local_dir) / filename
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_text("full")
        return str(dest)

    monkeypatch.setenv("KINK_HF_REQUIRE_FULL_CATALOG", "1")
    monkeypatch.setenv("KINK_HF_DATASET_REPO", "dummy/dataset")
    monkeypatch.setenv("KINK_HF_DATASET_FILENAME", "store_slim.db")
    with patch("huggingface_hub.hf_hub_download", side_effect=fake_download):
        out = hf_bootstrap.ensure_store_db(target)
    assert out == target.resolve()
    assert target.read_text() == "full"


def test_ensure_store_db_require_full_uses_default_repo_when_unset(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
    """Space UI can clear KINK_HF_DATASET_REPO; bootstrap still targets the default dataset."""
    from backend import hf_bootstrap

    target = tmp_path / "kink_store.db"
    seen: dict[str, str] = {}

    def fake_download(*, repo_id, filename, revision, token, local_dir, repo_type, **kwargs):
        seen["repo_id"] = repo_id
        seen["filename"] = filename
        assert repo_type == "dataset"
        dest = Path(local_dir) / filename
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_text("seed")
        return str(dest)

    monkeypatch.setenv("KINK_HF_REQUIRE_FULL_CATALOG", "1")
    monkeypatch.delenv("KINK_HF_DATASET_REPO", raising=False)
    monkeypatch.delenv("KINK_HF_DATASET_FILENAME", raising=False)
    with patch("huggingface_hub.hf_hub_download", side_effect=fake_download):
        out = hf_bootstrap.ensure_store_db(target)
    assert seen["repo_id"] == hf_bootstrap.DEFAULT_KINK_HF_DATASET_REPO
    assert seen["filename"] == hf_bootstrap.DEFAULT_KINK_HF_REMOTE_DB_NAME
    assert out == target.resolve()
    assert target.read_text() == "seed"


def test_ensure_store_db_catalog_url_skips_hub(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
    from backend import hf_bootstrap

    target = tmp_path / "store.db"
    called: list[str] = []

    def fake_http(url: str, dest: Path) -> None:
        called.append(url)
        dest.write_bytes(b"from-url")

    monkeypatch.setenv("KINK_CATALOG_URL", "https://example.invalid/store.db")
    monkeypatch.setenv("KINK_HF_DATASET_REPO", "dummy/dataset")
    with patch.object(hf_bootstrap, "_download_store_from_http", side_effect=fake_http):
        with patch("huggingface_hub.hf_hub_download") as hub_dl:
            out = hf_bootstrap.ensure_store_db(target)
    hub_dl.assert_not_called()
    assert called == ["https://example.invalid/store.db"]
    assert target.read_bytes() == b"from-url"
    assert out == target.resolve()


def test_ensure_store_db_downloads_when_repo_set(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
    from backend import hf_bootstrap

    target = tmp_path / "store_slim.db"

    def fake_download(*, repo_id, filename, revision, token, local_dir, repo_type, **kwargs):
        assert repo_type == "dataset"
        dest = Path(local_dir) / filename
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_text("seed")
        return str(dest)

    monkeypatch.setenv("KINK_HF_REQUIRE_FULL_CATALOG", "1")
    monkeypatch.setenv("KINK_HF_DATASET_REPO", "dummy/dataset")
    monkeypatch.setenv("KINK_HF_DATASET_FILENAME", "store_slim.db")
    with patch("huggingface_hub.hf_hub_download", side_effect=fake_download):
        out = hf_bootstrap.ensure_store_db(target)
    assert out == target.resolve()
    assert target.read_text() == "seed"


def test_ensure_store_db_bundled_ignores_stale_repo_without_opt_in(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
    """Bundled image + KINK_HF_REQUIRE_FULL_CATALOG=0 must not hf_hub_download when only a stale DATASET_REPO is set."""
    from backend import hf_bootstrap

    seed = tmp_path / "bundle.db"
    seed.write_bytes(b"bundled-seed")
    monkeypatch.setenv("KINK_SEED_PATH", str(seed))
    monkeypatch.setenv("KINK_HF_REQUIRE_FULL_CATALOG", "0")
    monkeypatch.setenv("KINK_HF_DATASET_REPO", "dummy/dataset")
    monkeypatch.delenv("KINK_HF_USE_HUB_DATASET", raising=False)
    target = tmp_path / "out.db"
    with patch("huggingface_hub.hf_hub_download") as hub_dl:
        out = hf_bootstrap.ensure_store_db(target)
    hub_dl.assert_not_called()
    assert out == target.resolve()
    assert target.read_bytes() == b"bundled-seed"


def test_ensure_store_db_b2_s3_skips_hub(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
    """B2 S3 download when B2_* env is set; no Hub call."""
    from backend import hf_bootstrap

    target = tmp_path / "store.db"
    fake_client = MagicMock()

    def download_file(bucket: str, key: str, path: str) -> None:
        assert bucket == "my-bucket"
        assert key == "custom/key.db"
        Path(path).write_bytes(b"b2-bytes")

    fake_client.download_file = MagicMock(side_effect=download_file)

    monkeypatch.delenv("KINK_CATALOG_URL", raising=False)
    monkeypatch.setenv("B2_KEY_ID", "kid")
    monkeypatch.setenv("B2_APPLICATION_KEY", "secret")
    monkeypatch.setenv("B2_BUCKET", "my-bucket")
    monkeypatch.setenv("B2_REGION", "us-west-004")
    monkeypatch.setenv("KINK_B2_OBJECT_KEY", "custom/key.db")
    monkeypatch.setenv("KINK_HF_DATASET_REPO", "dummy/dataset")

    with patch("boto3.client", return_value=fake_client):
        with patch("huggingface_hub.hf_hub_download") as hub_dl:
            out = hf_bootstrap.ensure_store_db(target)
    hub_dl.assert_not_called()
    fake_client.download_file.assert_called_once()
    assert target.read_bytes() == b"b2-bytes"
    assert out == target.resolve()