test: migrate reranker upload-file sessions and ORM models to SQLite (#40659)

This commit is contained in:
Asuka Minato
2026-08-13 05:19:19 +00:00
committed by GitHub
parent 2f969afaf2
commit ff46c3d7d7
@@ -14,7 +14,6 @@ Tests follow the Arrange-Act-Assert pattern for clarity.
from datetime import UTC, datetime
from operator import itemgetter
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch
import pytest
@@ -480,10 +479,28 @@ class TestRerankModelRunnerMultimodal(_UsesSQLiteSession):
assert len(result) == 1
assert result[0].metadata["score"] == 0.88
def test_fetch_multimodal_rerank_builds_docs_and_calls_text_rerank(self, rerank_runner):
def test_fetch_multimodal_rerank_builds_docs_and_calls_text_rerank(
self, rerank_runner: RerankModelRunner, sqlite_session: Session
):
upload_file = UploadFile(
tenant_id="00000000-0000-0000-0000-000000000001",
storage_type=StorageType.LOCAL,
key="image-key",
name="document.png",
size=10,
extension="png",
mime_type="image/png",
created_by_role=CreatorUserRole.ACCOUNT,
created_by="00000000-0000-0000-0000-000000000002",
created_at=datetime.now(UTC),
used=True,
)
upload_file.id = "00000000-0000-0000-0000-000000000003"
sqlite_session.add(upload_file)
sqlite_session.commit()
image_doc = Document(
page_content="image-content",
metadata={"doc_id": "img-1", "doc_type": DocType.IMAGE},
metadata={"doc_id": upload_file.id, "doc_type": DocType.IMAGE},
provider="dify",
)
text_doc = Document(
@@ -499,7 +516,6 @@ class TestRerankModelRunnerMultimodal(_UsesSQLiteSession):
rerank_result = RerankResult(model="rerank-model", docs=[])
with (
patch.object(rerank_runner._session, "get", return_value=SimpleNamespace(key="image-key")),
patch("core.rag.rerank.rerank_model.storage.load_once", return_value=b"image-bytes") as mock_load_once,
patch.object(
rerank_runner,
@@ -527,14 +543,11 @@ class TestRerankModelRunnerMultimodal(_UsesSQLiteSession):
)
rerank_result = RerankResult(model="rerank-model", docs=[])
with (
patch.object(rerank_runner._session, "get", return_value=None),
patch.object(
rerank_runner,
"fetch_text_rerank",
return_value=(rerank_result, [image_doc]),
) as mock_text_rerank,
):
with patch.object(
rerank_runner,
"fetch_text_rerank",
return_value=(rerank_result, [image_doc]),
) as mock_text_rerank:
result, unique_documents = rerank_runner.fetch_multimodal_rerank(
query="python",
documents=[image_doc],
@@ -596,13 +609,12 @@ class TestRerankModelRunnerMultimodal(_UsesSQLiteSession):
assert "user" not in invoke_kwargs
def test_fetch_multimodal_rerank_raises_when_query_image_not_found(self, rerank_runner):
with patch.object(rerank_runner._session, "get", return_value=None):
with pytest.raises(ValueError, match="Upload file not found for query"):
rerank_runner.fetch_multimodal_rerank(
query="missing-upload-id",
documents=[],
query_type=QueryType.IMAGE_QUERY,
)
with pytest.raises(ValueError, match="Upload file not found for query"):
rerank_runner.fetch_multimodal_rerank(
query="missing-upload-id",
documents=[],
query_type=QueryType.IMAGE_QUERY,
)
def test_fetch_multimodal_rerank_rejects_unsupported_query_type(self, rerank_runner):
with pytest.raises(ValueError, match="is not supported"):