fix: mini fixes

tests/scripts/test_prefetch_models.py

@@ -0,0 +1,96 @@
+"""Model pre-fetching, which the Docker build runs to bake artifacts in."""
+
+import sys
+import types
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from application.scripts import prefetch_models
+from application.vectorstore.model_registry import GRANITE_311M, MPNET, OPENAI_ADA_002
+
+
+@pytest.fixture
+def fake_fastembed():
+    """Stand in for FastEmbed so nothing is downloaded."""
+    text_embedding = MagicMock()
+    pooling = types.SimpleNamespace(CLS="CLS", MEAN="MEAN")
+    module = types.ModuleType("fastembed")
+    module.TextEmbedding = text_embedding
+    desc = types.ModuleType("fastembed.common.model_description")
+    desc.PoolingType = pooling
+    desc.ModelSource = lambda hf=None: {"hf": hf}
+    with patch.dict(
+        sys.modules,
+        {"fastembed": module, "fastembed.common.model_description": desc},
+    ):
+        yield text_embedding
+
+
+class TestPrefetch:
+    def test_defaults_cover_legacy_and_new_install(self):
+        """An upgraded image must still serve mpnet; a new one needs granite."""
+        assert MPNET.name in prefetch_models.DEFAULT_MODELS
+        assert GRANITE_311M.name in prefetch_models.DEFAULT_MODELS
+
+    def test_fetches_named_models(self, fake_fastembed):
+        fetched = prefetch_models.prefetch([GRANITE_311M.name])
+        assert fetched == [GRANITE_311M.repo]
+        assert fake_fastembed.call_args.kwargs["model_name"] == GRANITE_311M.repo
+
+    def test_registers_with_registry_pooling_and_dim(self, fake_fastembed):
+        prefetch_models.prefetch([GRANITE_311M.name])
+        kwargs = fake_fastembed.add_custom_model.call_args.kwargs
+        assert kwargs["model"] == GRANITE_311M.repo
+        assert kwargs["pooling"] == "CLS"
+        assert kwargs["dim"] == GRANITE_311M.dimension
+        assert kwargs["model_file"] == GRANITE_311M.onnx_file
+
+    def test_mean_pooled_model_registers_as_mean(self, fake_fastembed):
+        prefetch_models.prefetch([MPNET.name])
+        assert fake_fastembed.add_custom_model.call_args.kwargs["pooling"] == "MEAN"
+
+    def test_aliases_resolve(self, fake_fastembed):
+        assert prefetch_models.prefetch(["granite-311m"]) == [GRANITE_311M.repo]
+
+    def test_cache_dir_is_forwarded(self, fake_fastembed):
+        prefetch_models.prefetch([MPNET.name], cache_dir="/models")
+        assert fake_fastembed.call_args.kwargs["cache_dir"] == "/models"
+
+    def test_cache_dir_omitted_when_absent(self, fake_fastembed):
+        prefetch_models.prefetch([MPNET.name])
+        assert "cache_dir" not in fake_fastembed.call_args.kwargs
+
+    def test_remote_only_model_is_skipped(self, fake_fastembed):
+        """OpenAI embeddings have no local artifacts to cache."""
+        assert prefetch_models.prefetch([OPENAI_ADA_002.name]) == []
+        fake_fastembed.assert_not_called()
+
+    def test_unknown_model_fails_loudly(self, fake_fastembed):
+        """A silent skip at build time is a download at run time, offline."""
+        with pytest.raises(SystemExit) as excinfo:
+            prefetch_models.prefetch(["nope/nope"])
+        assert "nope/nope" in str(excinfo.value)
+        assert MPNET.name in str(excinfo.value)
+
+    def test_several_models_in_one_run(self, fake_fastembed):
+        fetched = prefetch_models.prefetch([MPNET.name, GRANITE_311M.name])
+        assert fetched == [MPNET.repo, GRANITE_311M.repo]
+
+
+class TestMain:
+    def test_no_args_fetches_the_defaults(self, fake_fastembed):
+        with patch.object(prefetch_models, "prefetch", return_value=[]) as spy:
+            assert prefetch_models.main([]) == 0
+        assert spy.call_args.args[0] == list(prefetch_models.DEFAULT_MODELS)
+
+    def test_explicit_args_override_the_defaults(self, fake_fastembed):
+        with patch.object(prefetch_models, "prefetch", return_value=[]) as spy:
+            prefetch_models.main(["granite-97m"])
+        assert spy.call_args.args[0] == ["granite-97m"]
+
+    def test_cache_dir_read_from_environment(self, fake_fastembed, monkeypatch):
+        monkeypatch.setenv("EMBEDDINGS_CACHE_DIR", "/app/models")
+        with patch.object(prefetch_models, "prefetch", return_value=[]) as spy:
+            prefetch_models.main(["granite-97m"])
+        assert spy.call_args.args[1] == "/app/models"

tests/scripts/test_reembed.py

@@ -122,3 +122,112 @@ class TestFaissRebuild:
             reembed.VectorCreator, "create_vectorstore", return_value=existing
         ):
             assert reembed.reembed_faiss("s1", batch_size=8, dry_run=False) == (0, 0)
+
+
+class TestPgvectorWithoutTheExtension:
+    """Mocked pgvector paths.
+
+    The live tests in ``test_reembed_pgvector_live`` skip wherever the cluster
+    has no pgvector build -- which includes CI -- so the SQL shape and the
+    batching contract are pinned here too.
+    """
+
+    @pytest.fixture
+    def store(self):
+        store = MagicMock()
+        store._table_name = "documents"
+        store._vector_column = "embedding"
+        cursor = MagicMock()
+        cursor.fetchall.return_value = [(1, "alpha"), (2, "beta"), (3, "gamma")]
+        conn = MagicMock()
+        conn.cursor.return_value = cursor
+        store._get_connection.return_value = conn
+        store._embedding.embed_documents.side_effect = lambda texts: [
+            [0.5] * 4 for _ in texts
+        ]
+        with patch.object(
+            reembed.VectorCreator, "create_vectorstore", return_value=store
+        ):
+            yield store, conn, cursor
+
+    def test_reads_and_rewrites_every_chunk(self, store):
+        _, conn, cursor = store
+        seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
+        assert (seen, written) == (3, 3)
+        cursor.executemany.assert_called_once()
+        conn.commit.assert_called()
+
+    def test_dry_run_neither_embeds_nor_writes(self, store):
+        fake_store, conn, cursor = store
+        seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=True)
+        assert (seen, written) == (3, 0)
+        fake_store._embedding.embed_documents.assert_not_called()
+        cursor.executemany.assert_not_called()
+
+    def test_batches_commit_separately(self, store):
+        _, conn, cursor = store
+        reembed.reembed_pgvector("s1", batch_size=2, dry_run=False)
+        # 3 rows at batch 2 is two write transactions.
+        assert cursor.executemany.call_count == 2
+        assert conn.commit.call_count == 2
+
+    def test_failed_batch_rolls_back_and_raises(self, store):
+        fake_store, conn, cursor = store
+        cursor.executemany.side_effect = RuntimeError("write failed")
+        with pytest.raises(RuntimeError):
+            reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
+        conn.rollback.assert_called_once()
+
+    def test_connection_is_returned_even_on_failure(self, store):
+        fake_store, _, cursor = store
+        cursor.executemany.side_effect = RuntimeError("boom")
+        with pytest.raises(RuntimeError):
+            reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
+        fake_store.close.assert_called_once()
+
+    def test_null_text_does_not_crash_the_embed_call(self, store):
+        fake_store, _, cursor = store
+        cursor.fetchall.return_value = [(1, None), (2, "beta")]
+        seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
+        assert (seen, written) == (2, 2)
+        assert fake_store._embedding.embed_documents.call_args.args[0] == ["", "beta"]
+
+    def test_empty_source_is_a_no_op(self, store):
+        fake_store, _, cursor = store
+        cursor.fetchall.return_value = []
+        assert reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) == (0, 0)
+        fake_store._embedding.embed_documents.assert_not_called()
+
+    def test_source_discovery_returns_sorted_ids(self, store):
+        _, _, cursor = store
+        cursor.fetchall.return_value = [("b",), ("a",)]
+        assert reembed.list_source_ids("pgvector") == ["b", "a"]
+
+
+class TestFaissSourceDiscovery:
+    def test_source_ids_come_from_index_directories(self):
+        storage = MagicMock()
+        storage.list_files.return_value = [
+            "indexes/src-a/index.faiss",
+            "indexes/src-a/index.pkl",
+            "indexes/src-b/index.faiss",
+        ]
+        with patch(
+            "application.storage.storage_creator.StorageCreator.get_storage",
+            return_value=storage,
+        ):
+            assert reembed.list_source_ids("faiss") == ["src-a", "src-b"]
+
+    def test_storage_failure_is_reported_as_a_usable_error(self):
+        storage = MagicMock()
+        storage.list_files.side_effect = OSError("permission denied")
+        with patch(
+            "application.storage.storage_creator.StorageCreator.get_storage",
+            return_value=storage,
+        ):
+            with pytest.raises(reembed.ReembedError, match="permission denied"):
+                reembed.list_source_ids("faiss")
+
+    def test_unsupported_store_error_names_the_alternatives(self):
+        with patch.object(reembed.settings, "VECTOR_STORE", "milvus", create=True):
+            assert reembed.main([]) == 2

tests/scripts/test_reembed_pgvector_live.py

@@ -44,11 +44,15 @@ def _dsn(info) -> str:
 
 @pytest.fixture(autouse=True)
 def _close_pools():
+    """Never leak a pool into another test; the DSN dies with the test DB."""
     yield
     for dsn, pool in list(pgvector_module._POOLS.items()):
         try:
             pool.close()
         except Exception:
+            # Teardown only: the ephemeral cluster may already be gone, and a
+            # failure to close a pool for a dead DSN must not fail the test
+            # that just passed. Dropping the entry below is what matters.
             pass
         pgvector_module._POOLS.pop(dsn, None)