opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-30 00:39:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: use SQLite sessions in core rag (#39106)

Browse Source
This commit is contained in:
Asuka Minato 2026-07-29 14:52:45 +09:00 committed by GitHub
parent c5b0772ab0
commit 0290c8d525
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 79 additions and 85 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

164
api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py
View File

@ -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)]