mirror of
https://github.com/langgenius/dify.git
synced 2026-10-30 18:34:29 +08:00
test: migrate reranker upload-file sessions and ORM models to SQLite (#40659)
This commit is contained in:
@@ -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"):
|
||||
|
||||
Reference in New Issue
Block a user