From 0290c8d5253da1262387334271bb01078d3d2a9c Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Wed, 29 Jul 2026 14:52:45 +0900 Subject: [PATCH] test: use SQLite sessions in core rag (#39106) --- .../core/rag/extractor/test_pdf_extractor.py | 164 +++++++++--------- 1 file changed, 79 insertions(+), 85 deletions(-) diff --git a/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py index c41a0752b48..3f5cf0d37cb 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py @@ -1,83 +1,68 @@ -from types import SimpleNamespace +from dataclasses import dataclass from unittest.mock import MagicMock, patch +from uuid import uuid4 import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session import core.rag.extractor.pdf_extractor as pe +from models.model import UploadFile + +TENANT_ID = str(uuid4()) +USER_ID = str(uuid4()) + + +class _Storage: + saves: list[tuple[str, bytes]] + + def __init__(self) -> None: + self.saves = [] + + def save(self, key: str, data: bytes) -> None: + self.saves.append((key, data)) + + +class _DatabaseBinding: + session: Session + + def __init__(self, session: Session) -> None: + self.session = session + + +@dataclass(frozen=True) +class _Dependencies: + storage: _Storage + session: Session @pytest.fixture -def mock_dependencies(monkeypatch: pytest.MonkeyPatch): - # Mock storage - saves = [] - - def save(key, data): - saves.append((key, data)) - - monkeypatch.setattr(pe, "storage", SimpleNamespace(save=save)) - - # Mock db - class DummySession: - def __init__(self): - self.added = [] - self.committed = False - - def add(self, obj): - self.added.append(obj) - - def add_all(self, objs): - self.added.extend(objs) - - def commit(self): - self.committed = True - - db_stub = SimpleNamespace(session=DummySession()) - monkeypatch.setattr(pe, "db", db_stub) - - # Mock UploadFile - class FakeUploadFile: - DEFAULT_ID = "test_file_id" - - def __init__(self, **kwargs): - # Assign id from DEFAULT_ID, allow override via kwargs if needed - self.id = self.DEFAULT_ID - for k, v in kwargs.items(): - setattr(self, k, v) - - monkeypatch.setattr(pe, "UploadFile", FakeUploadFile) - - # Mock config +def mock_dependencies(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> _Dependencies: + storage = _Storage() + monkeypatch.setattr(pe, "storage", storage) + monkeypatch.setattr(pe, "db", _DatabaseBinding(sqlite_session)) monkeypatch.setattr(pe.dify_config, "FILES_URL", "http://files.local") monkeypatch.setattr(pe.dify_config, "INTERNAL_FILES_URL", None) monkeypatch.setattr(pe.dify_config, "STORAGE_TYPE", "local") - - return SimpleNamespace(saves=saves, db=db_stub, UploadFile=FakeUploadFile) + return _Dependencies(storage=storage, session=sqlite_session) @pytest.mark.parametrize( - ("image_bytes", "expected_mime", "expected_ext", "file_id"), + ("image_bytes", "expected_mime", "expected_ext"), [ - (b"\xff\xd8\xff some jpeg", "image/jpeg", "jpg", "test_file_id_jpeg"), - (b"\x89PNG\r\n\x1a\n some png", "image/png", "png", "test_file_id_png"), + (b"\xff\xd8\xff some jpeg", "image/jpeg", "jpg"), + (b"\x89PNG\r\n\x1a\n some png", "image/png", "png"), ], ) +@pytest.mark.parametrize("sqlite_session", [(UploadFile,)], indirect=True) @pytest.mark.parametrize("inject_session", [False, True]) def test_extract_images_formats( - mock_dependencies, - monkeypatch: pytest.MonkeyPatch, - image_bytes, - expected_mime, - expected_ext, - file_id, + mock_dependencies: _Dependencies, + image_bytes: bytes, + expected_mime: str, + expected_ext: str, inject_session: bool, ): - saves = mock_dependencies.saves - db_stub = mock_dependencies.db - - # Customize FakeUploadFile id for this test case. - # Using monkeypatch ensures the class attribute is reset between parameter sets. - monkeypatch.setattr(mock_dependencies.UploadFile, "DEFAULT_ID", file_id) - # Mock page and image objects mock_page = MagicMock() mock_image_obj = MagicMock() @@ -91,25 +76,32 @@ def test_extract_images_formats( extractor = pe.PdfExtractor( file_path="test.pdf", - tenant_id="t1", - user_id="u1", - session=db_stub.session if inject_session else None, + tenant_id=TENANT_ID, + user_id=USER_ID, + session=mock_dependencies.session if inject_session else None, ) # We need to handle the import inside _extract_images - with patch("pypdfium2.raw", autospec=True) as mock_raw: + with ( + patch("pypdfium2.raw", autospec=True) as mock_raw, + patch.object( + mock_dependencies.session, + "commit", + wraps=mock_dependencies.session.commit, + ) as commit, + ): mock_raw.FPDF_PAGEOBJ_IMAGE = 1 result = extractor._extract_images(mock_page) - assert f"![image](http://files.local/files/{file_id}/file-preview)" in result - assert len(saves) == 1 - assert saves[0][1] == image_bytes - assert len(db_stub.session.added) == 1 - assert db_stub.session.added[0].tenant_id == "t1" - assert db_stub.session.added[0].size == len(image_bytes) - assert db_stub.session.added[0].mime_type == expected_mime - assert db_stub.session.added[0].extension == expected_ext - assert db_stub.session.committed is not inject_session + assert commit.called is not inject_session + upload_file = mock_dependencies.session.scalar(select(UploadFile)) + assert upload_file is not None + assert f"![image](http://files.local/files/{upload_file.id}/file-preview)" in result + assert mock_dependencies.storage.saves == [(upload_file.key, image_bytes)] + assert upload_file.tenant_id == TENANT_ID + assert upload_file.size == len(image_bytes) + assert upload_file.mime_type == expected_mime + assert upload_file.extension == expected_ext @pytest.mark.parametrize( @@ -120,14 +112,17 @@ def test_extract_images_formats( (Exception("Failed to get objects"), None), # Exception raised ], ) -def test_extract_images_get_objects_scenarios(mock_dependencies, get_objects_side_effect, get_objects_return_value): +@pytest.mark.parametrize("sqlite_session", [(UploadFile,)], indirect=True) +def test_extract_images_get_objects_scenarios( + mock_dependencies: _Dependencies, get_objects_side_effect, get_objects_return_value +): mock_page = MagicMock() if get_objects_side_effect: mock_page.get_objects.side_effect = get_objects_side_effect else: mock_page.get_objects.return_value = get_objects_return_value - extractor = pe.PdfExtractor(file_path="test.pdf", tenant_id="t1", user_id="u1") + extractor = pe.PdfExtractor(file_path="test.pdf", tenant_id=TENANT_ID, user_id=USER_ID) with patch("pypdfium2.raw", autospec=True) as mock_raw: mock_raw.FPDF_PAGEOBJ_IMAGE = 1 @@ -136,7 +131,8 @@ def test_extract_images_get_objects_scenarios(mock_dependencies, get_objects_sid assert result == "" -def test_extract_calls_extract_images(mock_dependencies, monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("sqlite_session", [(UploadFile,)], indirect=True) +def test_extract_calls_extract_images(mock_dependencies: _Dependencies, monkeypatch: pytest.MonkeyPatch): # Mock pypdfium2 mock_pdf_doc = MagicMock() mock_page = MagicMock() @@ -152,7 +148,7 @@ def test_extract_calls_extract_images(mock_dependencies, monkeypatch: pytest.Mon mock_blob = MagicMock() mock_blob.source = "test.pdf" with patch("core.rag.extractor.pdf_extractor.Blob.from_path", return_value=mock_blob, autospec=True): - extractor = pe.PdfExtractor(file_path="test.pdf", tenant_id="t1", user_id="u1") + extractor = pe.PdfExtractor(file_path="test.pdf", tenant_id=TENANT_ID, user_id=USER_ID) # Mock _extract_images to return a known string monkeypatch.setattr(extractor, "_extract_images", lambda p: "![image](img_url)") @@ -165,10 +161,8 @@ def test_extract_calls_extract_images(mock_dependencies, monkeypatch: pytest.Mon assert documents[0].metadata["page"] == 0 -def test_extract_images_failures(mock_dependencies): - saves = mock_dependencies.saves - db_stub = mock_dependencies.db - +@pytest.mark.parametrize("sqlite_session", [(UploadFile,)], indirect=True) +def test_extract_images_failures(mock_dependencies: _Dependencies): # Mock page and image objects mock_page = MagicMock() mock_image_obj_fail = MagicMock() @@ -187,14 +181,14 @@ def test_extract_images_failures(mock_dependencies): mock_page.get_objects.return_value = [mock_image_obj_fail, mock_image_obj_ok] - extractor = pe.PdfExtractor(file_path="test.pdf", tenant_id="t1", user_id="u1") + extractor = pe.PdfExtractor(file_path="test.pdf", tenant_id=TENANT_ID, user_id=USER_ID) with patch("pypdfium2.raw", autospec=True) as mock_raw: mock_raw.FPDF_PAGEOBJ_IMAGE = 1 result = extractor._extract_images(mock_page) # Should have one success - assert "![image](http://files.local/files/test_file_id/file-preview)" in result - assert len(saves) == 1 - assert saves[0][1] == jpeg_bytes - assert db_stub.session.committed is True + upload_file = mock_dependencies.session.scalar(select(UploadFile)) + assert upload_file is not None + assert f"![image](http://files.local/files/{upload_file.id}/file-preview)" in result + assert mock_dependencies.storage.saves == [(upload_file.key, jpeg_bytes)]