backend/tests/test_embed_breed_single_pass.py

"""The storage pipeline must produce a picture's embedding AND breed from ONE model pass when the
embedder and breed classifier are the same HF model (spec §9.1). Also covers the mock fallback."""
import numpy as np
from sqlalchemy import select

from app.db import SessionLocal
from app.models import BreedPrediction, Embedding, KnownDog, User
from app.models.base import SubjectType, UserRole
from app.security import hash_password
from app.services import images as images_svc
from app.services.images import embed_and_breed_picture, process_and_store_picture


def _img(seed):
    from scripts.make_sample_images import make_image

    return make_image(seed)


def _make_known_picture(db):
    user = User(name="O", email="o@example.com", zip="20001",
                password_hash=hash_password("password123"), role=UserRole.owner)
    db.add(user)
    db.flush()
    dog = KnownDog(owner_id=user.id, name="Rex", description="")
    db.add(dog)
    db.flush()
    # Store the image only (no embed/breed yet) so we can drive the generators explicitly.
    pic = process_and_store_picture(
        db, subject_type=SubjectType.known, subject_id=dog.id, data=_img(101),
        generate_embedding=False, generate_breed=False,
    )
    db.commit()
    return pic


def test_mock_fallback_writes_both_and_is_idempotent():
    db = SessionLocal()
    try:
        pic = _make_known_picture(db)
        emb, breed = embed_and_breed_picture(db, pic, skip_if_exists=True)
        db.commit()
        assert emb and breed
        assert db.execute(
            select(Embedding).where(Embedding.picture_id == pic.id)
        ).first() is not None
        assert db.execute(
            select(BreedPrediction).where(BreedPrediction.picture_id == pic.id)
        ).first() is not None
        # Re-run: nothing new to write.
        again = embed_and_breed_picture(db, pic, skip_if_exists=True)
        assert again == (False, False)
    finally:
        db.close()


class _FakeHFEmbedder:
    """Stands in for HFEmbedder: records how many times the model 'forward' runs."""
    name = "hf-embed"
    version = "Fake-Model"
    dim = 4
    calls = 0

    def embed_and_breed(self, paths, top_k):
        type(self).calls += 1  # one call == one forward pass over the batch
        vec = np.ones(self.dim, dtype=np.float32) / 2.0
        labels = [("border collie", 0.9), ("kelpie", 0.1)][:top_k]
        return [(vec, labels) for _ in paths]


def test_same_hf_model_runs_the_model_once_for_both(monkeypatch):
    db = SessionLocal()
    try:
        pic = _make_known_picture(db)

        # Pretend embedder + breed are the SAME HF model.
        monkeypatch.setattr(images_svc.settings, "embedder", "hf", raising=False)
        monkeypatch.setattr(images_svc.settings, "breed_classifier", "hf", raising=False)
        monkeypatch.setattr(images_svc.settings, "embedder_hf_model", "x/Fake-Model", raising=False)
        monkeypatch.setattr(images_svc.settings, "breed_model", "x/Fake-Model", raising=False)
        monkeypatch.setattr(images_svc.settings, "breed_top_k", 2, raising=False)
        fake = _FakeHFEmbedder()
        monkeypatch.setattr(images_svc, "get_embedder", lambda: fake)

        emb, breed = embed_and_breed_picture(db, pic, skip_if_exists=False)
        db.commit()

        assert emb and breed
        assert _FakeHFEmbedder.calls == 1  # the model ran exactly once, not twice

        # Embedding tagged hf-embed, breed rows tagged hf-breed/Fake-Model (shared naming).
        e = db.execute(
            select(Embedding).where(
                Embedding.picture_id == pic.id, Embedding.model_name == "hf-embed"
            )
        ).scalar_one()
        assert e.dim == 4
        breeds = db.execute(
            select(BreedPrediction).where(
                BreedPrediction.picture_id == pic.id,
                BreedPrediction.model_name == "hf-breed",
                BreedPrediction.model_version == "Fake-Model",
            )
        ).scalars().all()
        assert [b.label for b in breeds] == ["border collie", "kelpie"]
    finally:
        db.close()