From 3453fdce44c6409268efb8b4089f2fecf7c132d7 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Thu, 30 Jul 2026 12:04:18 +0900 Subject: [PATCH] test: use SQLite sessions in core rag (#39092) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../rag/embedding/test_cached_embedding.py | 196 +++++++++++------- 1 file changed, 124 insertions(+), 72 deletions(-) diff --git a/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py b/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py index b33a7ba725c..97b1327bc88 100644 --- a/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py +++ b/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py @@ -8,19 +8,55 @@ This test file covers the methods not fully tested in test_embedding_service.py: import base64 import logging +from dataclasses import dataclass from decimal import Decimal from unittest.mock import Mock, patch import numpy as np import pytest -from sqlalchemy.exc import IntegrityError +from sqlalchemy import event, func, select +from sqlalchemy.orm import Session +from core.rag.embedding import cached_embedding as cached_embedding_module from core.rag.embedding.cached_embedding import CacheEmbedding from graphon.model_runtime.entities.model_entities import ModelPropertyKey from graphon.model_runtime.entities.text_embedding_entities import EmbeddingResult, EmbeddingUsage from models.dataset import Embedding +@dataclass(frozen=True) +class _DatabaseBinding: + session: Session + + +@pytest.fixture +def embedding_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: + """Bind CacheEmbedding to the shared SQLite session.""" + + monkeypatch.setattr(cached_embedding_module, "db", _DatabaseBinding(session=sqlite_session)) + return sqlite_session + + +def _persist_embedding( + session: Session, + *, + cache_key: str, + vector: list[float], + model_name: str = "vision-embedding-model", + provider_name: str = "openai", +) -> Embedding: + embedding = Embedding( + model_name=model_name, + hash=cache_key, + provider_name=provider_name, + embedding=b"placeholder", + ) + embedding.set_embedding(vector) + session.add(embedding) + session.commit() + return embedding + + class TestCacheEmbeddingMultimodalDocuments: """Test suite for CacheEmbedding.embed_multimodal_documents method.""" @@ -65,27 +101,29 @@ class TestCacheEmbeddingMultimodalDocuments: ) def test_embed_single_multimodal_document_cache_miss( - self, mock_model_instance, sample_multimodal_result: EmbeddingResult + self, + mock_model_instance, + sample_multimodal_result: EmbeddingResult, + embedding_session: Session, ): """Test embedding a single multimodal document when cache is empty.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [{"file_id": "file123", "content": "test content"}] - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result + mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result - result = cache_embedding.embed_multimodal_documents(documents) + result = cache_embedding.embed_multimodal_documents(documents) - assert len(result) == 1 - assert isinstance(result[0], list) - assert len(result[0]) == 1536 + assert len(result) == 1 + assert isinstance(result[0], list) + assert len(result[0]) == 1536 - mock_model_instance.invoke_multimodal_embedding.assert_called_once() - mock_session.add.assert_called_once() - mock_session.commit.assert_called_once() + mock_model_instance.invoke_multimodal_embedding.assert_called_once() + persisted = embedding_session.scalar(select(Embedding).where(Embedding.hash == "file123")) + assert persisted is not None + assert persisted.get_embedding() == result[0] - def test_embed_multiple_multimodal_documents_cache_miss(self, mock_model_instance): + def test_embed_multiple_multimodal_documents_cache_miss(self, mock_model_instance, embedding_session: Session): """Test embedding multiple multimodal documents when cache is empty.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [ @@ -116,16 +154,15 @@ class TestCacheEmbeddingMultimodalDocuments: usage=usage, ) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result + mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - result = cache_embedding.embed_multimodal_documents(documents) + result = cache_embedding.embed_multimodal_documents(documents) - assert len(result) == 3 - assert all(len(emb) == 1536 for emb in result) + assert len(result) == 3 + assert all(len(emb) == 1536 for emb in result) + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3 - def test_embed_multimodal_documents_cache_hit(self, mock_model_instance): + def test_embed_multimodal_documents_cache_hit(self, mock_model_instance, embedding_session: Session): """Test embedding multimodal documents when embeddings are cached.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [{"file_id": "file123"}] @@ -133,19 +170,15 @@ class TestCacheEmbeddingMultimodalDocuments: cached_vector = np.random.randn(1536) normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - mock_cached_embedding = Mock(spec=Embedding) - mock_cached_embedding.get_embedding.return_value = normalized_cached + _persist_embedding(embedding_session, cache_key="file123", vector=normalized_cached) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = mock_cached_embedding + result = cache_embedding.embed_multimodal_documents(documents) - result = cache_embedding.embed_multimodal_documents(documents) + assert len(result) == 1 + assert result[0] == normalized_cached + mock_model_instance.invoke_multimodal_embedding.assert_not_called() - assert len(result) == 1 - assert result[0] == normalized_cached - mock_model_instance.invoke_multimodal_embedding.assert_not_called() - - def test_embed_multimodal_documents_partial_cache_hit(self, mock_model_instance): + def test_embed_multimodal_documents_partial_cache_hit(self, mock_model_instance, embedding_session: Session): """Test embedding multimodal documents with mixed cache hits and misses.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [ @@ -157,8 +190,7 @@ class TestCacheEmbeddingMultimodalDocuments: cached_vector = np.random.randn(1536) normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - mock_cached_embedding = Mock(spec=Embedding) - mock_cached_embedding.get_embedding.return_value = normalized_cached + _persist_embedding(embedding_session, cache_key="cached_file", vector=normalized_cached) new_embeddings = [] for _ in range(2): @@ -182,16 +214,20 @@ class TestCacheEmbeddingMultimodalDocuments: usage=usage, ) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.side_effect = [mock_cached_embedding, None, None] - mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result + mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - result = cache_embedding.embed_multimodal_documents(documents) + result = cache_embedding.embed_multimodal_documents(documents) - assert len(result) == 3 - assert result[0] == normalized_cached + assert len(result) == 3 + assert result[0] == normalized_cached + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3 - def test_embed_multimodal_documents_nan_handling(self, mock_model_instance, caplog: pytest.LogCaptureFixture): + def test_embed_multimodal_documents_nan_handling( + self, + mock_model_instance, + embedding_session: Session, + caplog: pytest.LogCaptureFixture, + ): """Test handling of NaN values in multimodal embeddings.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [{"file_id": "valid"}, {"file_id": "nan"}] @@ -215,20 +251,19 @@ class TestCacheEmbeddingMultimodalDocuments: usage=usage, ) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result + mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - with caplog.at_level(logging.WARNING, logger="core.rag.embedding.cached_embedding"): - result = cache_embedding.embed_multimodal_documents(documents) + with caplog.at_level(logging.WARNING, logger="core.rag.embedding.cached_embedding"): + result = cache_embedding.embed_multimodal_documents(documents) - assert len(result) == 2 - assert result[0] is not None - assert result[1] is None + assert len(result) == 2 + assert result[0] is not None + assert result[1] is None + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 1 - assert any(record.levelno == logging.WARNING for record in caplog.records) + assert any(record.levelno == logging.WARNING for record in caplog.records) - def test_embed_multimodal_documents_large_batch(self, mock_model_instance): + def test_embed_multimodal_documents_large_batch(self, mock_model_instance, embedding_session: Session): """Test embedding large batch of multimodal documents respecting MAX_CHUNKS.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [{"file_id": f"file{i}"} for i in range(25)] @@ -256,49 +291,66 @@ class TestCacheEmbeddingMultimodalDocuments: usage=usage, ) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None + batch_results = [create_batch_result(10), create_batch_result(10), create_batch_result(5)] + mock_model_instance.invoke_multimodal_embedding.side_effect = batch_results - batch_results = [create_batch_result(10), create_batch_result(10), create_batch_result(5)] - mock_model_instance.invoke_multimodal_embedding.side_effect = batch_results + result = cache_embedding.embed_multimodal_documents(documents) - result = cache_embedding.embed_multimodal_documents(documents) + assert len(result) == 25 + assert mock_model_instance.invoke_multimodal_embedding.call_count == 3 + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 25 - assert len(result) == 25 - assert mock_model_instance.invoke_multimodal_embedding.call_count == 3 - - def test_embed_multimodal_documents_api_error(self, mock_model_instance): + def test_embed_multimodal_documents_api_error(self, mock_model_instance, embedding_session: Session): """Test handling of API errors during multimodal embedding.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [{"file_id": "file123"}] - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_multimodal_embedding.side_effect = Exception("API Error") + mock_model_instance.invoke_multimodal_embedding.side_effect = Exception("API Error") - with pytest.raises(Exception) as exc_info: - cache_embedding.embed_multimodal_documents(documents) + with pytest.raises(Exception, match="API Error"): + cache_embedding.embed_multimodal_documents(documents) - assert "API Error" in str(exc_info.value) - mock_session.rollback.assert_called() + assert not embedding_session.in_transaction() + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0 def test_embed_multimodal_documents_integrity_error_during_transform( - self, mock_model_instance, sample_multimodal_result + self, + mock_model_instance, + sample_multimodal_result, + embedding_session: Session, ): """Test handling of IntegrityError during embedding transformation.""" cache_embedding = CacheEmbedding(mock_model_instance) documents = [{"file_id": "file123"}] - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result + mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result - mock_session.commit.side_effect = IntegrityError("Duplicate key", None, None) + injected = False + def add_competing_row(session: Session, _flush_context: object, _instances: object) -> None: + nonlocal injected + if injected: + return + pending = next(item for item in session.new if isinstance(item, Embedding)) + session.add( + Embedding( + model_name=pending.model_name, + hash=pending.hash, + provider_name=pending.provider_name, + embedding=pending.embedding, + ) + ) + injected = True + + event.listen(embedding_session, "before_flush", add_competing_row) + try: result = cache_embedding.embed_multimodal_documents(documents) + finally: + event.remove(embedding_session, "before_flush", add_competing_row) - assert len(result) == 1 - mock_session.rollback.assert_called() + assert len(result) == 1 + assert not embedding_session.in_transaction() + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0 class TestCacheEmbeddingMultimodalQuery: