tests/test_embed.py

"""Unit tests for the embedding-text builder (canlex/embed.py)."""
import unittest

from canlex.embed import embed_text, embed_texts, content_hash, plan


def _chunk(**kw):
    base = {"doc_type": "legislation", "act_short": "X", "marginal_note": "",
            "part": "", "heading": "", "text": ""}
    base.update(kw)
    return base


class EmbedTextTests(unittest.TestCase):
    def test_memorandum_title_comes_from_part(self):
        # A memo's marginal note is generic; its subject lives in 'part'.
        out = embed_text(_chunk(doc_type="memorandum", act_short="D-Memo",
                                marginal_note="Guidelines",
                                part="Value for Duty", text="body"))
        self.assertIn("Value for Duty", out)

    def test_legislation_title_comes_from_marginal_note(self):
        out = embed_text(_chunk(marginal_note="Application for protection",
                                part="PART 2", text="body"))
        self.assertIn("Application for protection", out)

    def test_title_is_repeated_for_emphasis(self):
        out = embed_text(_chunk(marginal_note="UNIQUEWORD", text="b"))
        self.assertEqual(out.count("UNIQUEWORD"), 2)


class IncrementalPlanTests(unittest.TestCase):
    """plan() decides which chunks the incremental build may reuse."""

    def _stored(self, chunk, rows=None):
        texts = embed_texts(chunk)
        return {chunk["id"]: (content_hash(texts),
                              rows if rows is not None else [[0.0]] * len(texts))}

    def test_unchanged_chunk_is_reused(self):
        c = _chunk(id="a", marginal_note="Title", text="body")
        reusable, to_embed = plan([c], self._stored(c))
        self.assertEqual([r[0] for r in reusable], ["a"])
        self.assertEqual(to_embed, [])

    def test_changed_text_is_re_embedded(self):
        old = _chunk(id="a", marginal_note="Title", text="old body")
        new = _chunk(id="a", marginal_note="Title", text="new body")
        reusable, to_embed = plan([new], self._stored(old))
        self.assertEqual(reusable, [])
        self.assertEqual([r[0] for r in to_embed], ["a"])

    def test_new_chunk_and_window_count_mismatch_are_re_embedded(self):
        c = _chunk(id="a", marginal_note="Title", text="body")
        reusable, to_embed = plan([c], {})
        self.assertEqual([r[0] for r in to_embed], ["a"])
        # same hash but a different number of stored rows: not trusted
        stored = self._stored(c, rows=[[0.0], [0.0]])
        reusable, to_embed = plan([c], stored)
        self.assertEqual([r[0] for r in to_embed], ["a"])

    def test_missing_hashes_means_nothing_reused(self):
        from canlex.embed import load_stored
        import pathlib
        self.assertEqual(load_stored(pathlib.Path("does-not-exist.npz")), {})


# ---- Audit follow-up 2026-09-21, group A items 1 and 2 ---------------------
def _embedder_cached():
    from huggingface_hub import try_to_load_from_cache
    from canlex.embed import EMB_REPO
    return all(isinstance(try_to_load_from_cache(EMB_REPO, name), str)
               for name in ("onnx/model_quantized.onnx", "tokenizer.json"))


@unittest.skipUnless(_embedder_cached(), "the embedder is not in the local model cache")
class CachedEmbedderTests(unittest.TestCase):
    """The query embedder resolves its files from the local cache with no
    request (rerank.hub_file), and the arena-shrinking RunOptions leave the
    vectors bit-identical."""

    def test_no_request_and_identical_vectors(self):
        from unittest import mock
        import numpy as np
        from canlex import rerank
        from canlex.embed import Embedder
        real = rerank.hf_hub_download
        calls = []

        def spy(*args, **kwargs):
            calls.append(kwargs.get("local_files_only", False))
            return real(*args, **kwargs)
        with mock.patch.object(rerank, "hf_hub_download", side_effect=spy):
            emb = Embedder()
        self.assertEqual(calls, [True, True])
        texts = ["Detention review within 48 hours.", "Seizure of currency."]
        with_shrink = emb.encode(texts)
        run = emb.session.run
        with mock.patch.object(emb.session, "run",
                               side_effect=lambda names, feed, *a: run(names, feed)):
            without = emb.encode(texts)
        self.assertTrue(np.array_equal(with_shrink, without))


if __name__ == "__main__":
    unittest.main()