mirror of
https://github.com/langgenius/dify.git
synced 2026-10-04 13:47:55 +08:00
test: use real pipeline and retriever fixtures (#43220)
This commit is contained in:
1 parent
efcf73b4da
commit
0f611811c2
4 files changed
+131
-51
No files matched your search
@@ -13,6 +13,7 @@ from sqlalchemy.orm import Session
|
||||
|
||||
from core.agent.errors import AgentMaxIterationError
|
||||
from core.agent.fc_agent_runner import FunctionCallAgentRunner
|
||||
from core.app.app_config.entities import DatasetRetrieveConfigEntity
|
||||
from core.app.apps.base_app_queue_manager import PublishFrom
|
||||
from core.app.entities.app_invoke_entities import CreditUsageCreatedBy
|
||||
from core.app.entities.queue_entities import QueueMessageFileEvent
|
||||
@@ -408,7 +409,31 @@ class TestBuildDatasetToolImageContents:
|
||||
"file "
|
||||
)
|
||||
|
||||
tool = MagicMock(spec=DatasetRetrieverTool)
|
||||
from core.tools.__base.tool_runtime import ToolRuntime
|
||||
from core.tools.entities.common_entities import I18nObject
|
||||
from core.tools.entities.tool_entities import ToolDescription, ToolEntity, ToolIdentity
|
||||
from core.tools.utils.dataset_retriever.dataset_retriever_tool import DatasetRetrieverTool as RetrievalTool
|
||||
|
||||
tool = DatasetRetrieverTool(
|
||||
entity=ToolEntity(
|
||||
identity=ToolIdentity(
|
||||
provider="dataset", author="Dify", name="dataset", label=I18nObject(en_US="Dataset")
|
||||
),
|
||||
parameters=[],
|
||||
description=ToolDescription(human=I18nObject(en_US="Retrieve dataset"), llm="Retrieve dataset"),
|
||||
),
|
||||
runtime=ToolRuntime(tenant_id=upload_file.tenant_id),
|
||||
retrieval_tool=RetrievalTool(
|
||||
tenant_id=upload_file.tenant_id,
|
||||
dataset_id="dataset-id",
|
||||
return_resource=True,
|
||||
retriever_from="dev",
|
||||
retrieve_config=DatasetRetrieveConfigEntity(
|
||||
retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE
|
||||
),
|
||||
inputs={},
|
||||
),
|
||||
)
|
||||
contents = runner._build_dataset_tool_image_contents(session, response, tool)
|
||||
|
||||
assert contents == [image_content]
|
||||
|
||||
@@ -2,14 +2,16 @@ import json
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from core.app.app_config.entities import WorkflowUIBasedAppConfig
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager
|
||||
from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager
|
||||
from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
|
||||
from core.app.entities.queue_entities import QueueWorkflowStartedEvent
|
||||
from core.workflow.system_variables import build_system_variables
|
||||
from extensions.ext_redis import RedisClientWrapper
|
||||
from graphon.entities import WorkflowStartReason
|
||||
from graphon.runtime import GraphRuntimeState
|
||||
from models.account import Account
|
||||
@@ -18,6 +20,11 @@ from models.workflow import Workflow, WorkflowType
|
||||
from tests.workflow_test_utils import build_test_variable_pool
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def bind_queue_redis(monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None:
|
||||
monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0])
|
||||
|
||||
|
||||
def _build_workflow_app_config() -> WorkflowUIBasedAppConfig:
|
||||
return WorkflowUIBasedAppConfig(
|
||||
tenant_id="tenant-id",
|
||||
@@ -46,8 +53,9 @@ def _build_runtime_state(run_id: str) -> GraphRuntimeState:
|
||||
|
||||
|
||||
def _build_pipeline(run_id: str, unbound_session_factory: sessionmaker[Session]) -> WorkflowAppGenerateTaskPipeline:
|
||||
queue_manager = MagicMock(spec=AppQueueManager)
|
||||
queue_manager.invoke_from = InvokeFrom.SERVICE_API
|
||||
queue_manager = WorkflowAppQueueManager(
|
||||
task_id="task-id", user_id="user-id", invoke_from=InvokeFrom.SERVICE_API, app_mode=AppMode.WORKFLOW
|
||||
)
|
||||
queue_manager.graph_runtime_state = _build_runtime_state(run_id)
|
||||
workflow = Workflow(
|
||||
id="workflow-id",
|
||||
|
||||
@@ -6,18 +6,28 @@ and the fallback used when a local message file references a missing upload.
|
||||
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from core.app.app_config.entities import (
|
||||
EasyUIBasedAppConfig,
|
||||
EasyUIBasedAppModelConfigFrom,
|
||||
ModelConfigEntity,
|
||||
PromptTemplateEntity,
|
||||
)
|
||||
from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
|
||||
from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom
|
||||
from core.app.entities.task_entities import MessageEndStreamResponse
|
||||
from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline
|
||||
from extensions.ext_redis import RedisClientWrapper
|
||||
from extensions.storage.storage_type import StorageType
|
||||
from graphon.file import FileTransferMethod, FileType
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import MessageFile, UploadFile
|
||||
from models.enums import ConversationFromSource, CreatorUserRole
|
||||
from models.model import AppMode, Conversation, Message, MessageFile, UploadFile
|
||||
from tests.unit_tests.core.model_fixtures import make_model_config
|
||||
|
||||
SQLITE_MODELS = (MessageFile, UploadFile)
|
||||
pytestmark = [
|
||||
@@ -37,19 +47,53 @@ class TestMessageEndStreamResponseFiles:
|
||||
monkeypatch.setattr("core.db.session_factory._session_maker", sqlite_session_maker)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_pipeline(self) -> Mock:
|
||||
"""Create the minimal pipeline collaborator required by the method under test."""
|
||||
|
||||
pipeline = Mock(spec=EasyUIBasedGenerateTaskPipeline)
|
||||
pipeline._message_id = str(uuid.uuid4())
|
||||
pipeline._task_state = Mock()
|
||||
pipeline._task_state.metadata = Mock()
|
||||
pipeline._task_state.metadata.model_dump = Mock(return_value={"test": "metadata"})
|
||||
pipeline._task_state.llm_result = Mock()
|
||||
pipeline._task_state.llm_result.usage = Mock()
|
||||
pipeline._application_generate_entity = Mock()
|
||||
pipeline._application_generate_entity.task_id = str(uuid.uuid4())
|
||||
return pipeline
|
||||
def pipeline(
|
||||
self, monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock]
|
||||
) -> EasyUIBasedGenerateTaskPipeline:
|
||||
"""Construct the pipeline with validated request, model, and task state."""
|
||||
monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0])
|
||||
request = ChatAppGenerateEntity(
|
||||
task_id=str(uuid.uuid4()),
|
||||
app_config=EasyUIBasedAppConfig(
|
||||
tenant_id="tenant-id",
|
||||
app_id="app-id",
|
||||
app_mode=AppMode.CHAT,
|
||||
app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG,
|
||||
app_model_config_dict={},
|
||||
model=ModelConfigEntity(provider="test-provider", model="test-model"),
|
||||
prompt_template=PromptTemplateEntity(prompt_type=PromptTemplateEntity.PromptType.SIMPLE),
|
||||
),
|
||||
model_conf=make_model_config(provider="test-provider", model="test-model", mode="chat"),
|
||||
inputs={},
|
||||
files=[],
|
||||
user_id="user-id",
|
||||
stream=True,
|
||||
invoke_from=InvokeFrom.WEB_APP,
|
||||
)
|
||||
conversation = Conversation(
|
||||
id=str(uuid.uuid4()),
|
||||
app_id="app-id",
|
||||
mode=AppMode.CHAT,
|
||||
name="Conversation",
|
||||
from_source=ConversationFromSource.API,
|
||||
inputs={},
|
||||
)
|
||||
message = Message(id=str(uuid.uuid4()), created_at=datetime.now())
|
||||
queue = MessageBasedAppQueueManager(
|
||||
task_id=request.task_id,
|
||||
user_id=request.user_id,
|
||||
invoke_from=request.invoke_from,
|
||||
conversation_id=conversation.id,
|
||||
app_mode=AppMode.CHAT,
|
||||
message_id=message.id,
|
||||
)
|
||||
return EasyUIBasedGenerateTaskPipeline(
|
||||
application_generate_entity=request,
|
||||
queue_manager=queue,
|
||||
conversation=conversation,
|
||||
message=message,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _message_file(
|
||||
@@ -120,7 +164,9 @@ class TestMessageEndStreamResponseFiles:
|
||||
session.add_all(rows)
|
||||
session.commit()
|
||||
|
||||
def test_message_end_with_no_files(self, sqlite_session: Session, mock_pipeline: Mock) -> None:
|
||||
def test_message_end_with_no_files(
|
||||
self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline
|
||||
) -> None:
|
||||
"""Rows for another message do not leak into an empty files array."""
|
||||
|
||||
unrelated_file = self._message_file(
|
||||
@@ -129,32 +175,35 @@ class TestMessageEndStreamResponseFiles:
|
||||
)
|
||||
self._persist(sqlite_session, unrelated_file)
|
||||
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert sqlite_session.get(MessageFile, unrelated_file.id) is unrelated_file
|
||||
assert isinstance(result, MessageEndStreamResponse)
|
||||
assert result.files == []
|
||||
assert result.id == mock_pipeline._message_id
|
||||
assert result.metadata == {"test": "metadata"}
|
||||
mock_pipeline._task_state.metadata.model_dump.assert_called_once_with(exclude_none=True)
|
||||
assert result.id == pipeline._message_id
|
||||
assert result.metadata == {
|
||||
"usage": pipeline._task_state.llm_result.usage.model_dump(exclude_none=True),
|
||||
"retriever_resources": [],
|
||||
"reasoning": {},
|
||||
}
|
||||
|
||||
def test_message_end_with_local_file(
|
||||
self,
|
||||
sqlite_session: Session,
|
||||
mock_pipeline: Mock,
|
||||
pipeline: EasyUIBasedGenerateTaskPipeline,
|
||||
message_file_local: MessageFile,
|
||||
upload_file: UploadFile,
|
||||
) -> None:
|
||||
"""Local files include persisted upload metadata and a signed URL."""
|
||||
|
||||
message_file_local.message_id = mock_pipeline._message_id
|
||||
message_file_local.message_id = pipeline._message_id
|
||||
self._persist(sqlite_session, message_file_local, upload_file)
|
||||
|
||||
with patch(
|
||||
"core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url",
|
||||
return_value="https://example.com/signed-url?signature=abc123",
|
||||
) as get_signed_url:
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert result.files is not None
|
||||
assert len(result.files) == 1
|
||||
@@ -172,14 +221,14 @@ class TestMessageEndStreamResponseFiles:
|
||||
get_signed_url.assert_called_once_with(upload_file_id=upload_file.id)
|
||||
|
||||
def test_message_end_with_remote_url(
|
||||
self, sqlite_session: Session, mock_pipeline: Mock, message_file_remote: MessageFile
|
||||
self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_remote: MessageFile
|
||||
) -> None:
|
||||
"""Remote files retain their source URL and derived filename."""
|
||||
|
||||
message_file_remote.message_id = mock_pipeline._message_id
|
||||
message_file_remote.message_id = pipeline._message_id
|
||||
self._persist(sqlite_session, message_file_remote)
|
||||
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert result.files is not None
|
||||
assert len(result.files) == 1
|
||||
@@ -194,15 +243,15 @@ class TestMessageEndStreamResponseFiles:
|
||||
assert file_dict["upload_file_id"] == message_file_remote.id
|
||||
|
||||
def test_message_end_with_tool_file_http(
|
||||
self, sqlite_session: Session, mock_pipeline: Mock, message_file_tool: MessageFile
|
||||
self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile
|
||||
) -> None:
|
||||
"""HTTP tool-file URLs pass through unchanged."""
|
||||
|
||||
message_file_tool.message_id = mock_pipeline._message_id
|
||||
message_file_tool.message_id = pipeline._message_id
|
||||
message_file_tool.url = "https://example.com/tool_file.png"
|
||||
self._persist(sqlite_session, message_file_tool)
|
||||
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert result.files is not None
|
||||
file_dict = result.files[0]
|
||||
@@ -212,18 +261,18 @@ class TestMessageEndStreamResponseFiles:
|
||||
assert file_dict["transfer_method"] == FileTransferMethod.TOOL_FILE.value
|
||||
|
||||
def test_message_end_with_tool_file_local(
|
||||
self, sqlite_session: Session, mock_pipeline: Mock, message_file_tool: MessageFile
|
||||
self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile
|
||||
) -> None:
|
||||
"""Local tool-file identifiers are signed at the external boundary."""
|
||||
|
||||
message_file_tool.message_id = mock_pipeline._message_id
|
||||
message_file_tool.message_id = pipeline._message_id
|
||||
self._persist(sqlite_session, message_file_tool)
|
||||
|
||||
with patch(
|
||||
"core.app.task_pipeline.message_file_utils.sign_tool_file",
|
||||
return_value="https://example.com/signed-tool-file.png?signature=xyz",
|
||||
) as sign_tool:
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert result.files is not None
|
||||
file_dict = result.files[0]
|
||||
@@ -234,11 +283,11 @@ class TestMessageEndStreamResponseFiles:
|
||||
sign_tool.assert_called_once_with(tool_file_id="tool_file_123", extension=".png")
|
||||
|
||||
def test_message_end_with_tool_file_long_extension(
|
||||
self, sqlite_session: Session, mock_pipeline: Mock, message_file_tool: MessageFile
|
||||
self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile
|
||||
) -> None:
|
||||
"""Overlong tool-file extensions use the safe binary fallback."""
|
||||
|
||||
message_file_tool.message_id = mock_pipeline._message_id
|
||||
message_file_tool.message_id = pipeline._message_id
|
||||
message_file_tool.url = "tool_file_abc.verylongextension"
|
||||
self._persist(sqlite_session, message_file_tool)
|
||||
|
||||
@@ -246,7 +295,7 @@ class TestMessageEndStreamResponseFiles:
|
||||
"core.app.task_pipeline.message_file_utils.sign_tool_file",
|
||||
return_value="https://example.com/signed.bin",
|
||||
) as sign_tool:
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert result.files is not None
|
||||
assert result.files[0]["extension"] == ".bin"
|
||||
@@ -255,39 +304,39 @@ class TestMessageEndStreamResponseFiles:
|
||||
def test_message_end_with_multiple_files(
|
||||
self,
|
||||
sqlite_session: Session,
|
||||
mock_pipeline: Mock,
|
||||
pipeline: EasyUIBasedGenerateTaskPipeline,
|
||||
message_file_local: MessageFile,
|
||||
message_file_remote: MessageFile,
|
||||
upload_file: UploadFile,
|
||||
) -> None:
|
||||
"""The response contains every persisted file associated with the message."""
|
||||
|
||||
message_file_local.message_id = mock_pipeline._message_id
|
||||
message_file_remote.message_id = mock_pipeline._message_id
|
||||
message_file_local.message_id = pipeline._message_id
|
||||
message_file_remote.message_id = pipeline._message_id
|
||||
self._persist(sqlite_session, message_file_local, message_file_remote, upload_file)
|
||||
|
||||
with patch(
|
||||
"core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url",
|
||||
return_value="https://example.com/signed-url?signature=abc123",
|
||||
):
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert result.files is not None
|
||||
assert {file["related_id"] for file in result.files} == {message_file_local.id, message_file_remote.id}
|
||||
|
||||
def test_message_end_with_local_file_no_upload_file(
|
||||
self, sqlite_session: Session, mock_pipeline: Mock, message_file_local: MessageFile
|
||||
self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_local: MessageFile
|
||||
) -> None:
|
||||
"""A missing upload row still signs the stored upload identifier."""
|
||||
|
||||
message_file_local.message_id = mock_pipeline._message_id
|
||||
message_file_local.message_id = pipeline._message_id
|
||||
self._persist(sqlite_session, message_file_local)
|
||||
|
||||
with patch(
|
||||
"core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url",
|
||||
return_value="https://example.com/fallback-url?signature=def456",
|
||||
) as get_signed_url:
|
||||
result = EasyUIBasedGenerateTaskPipeline._message_end_to_stream_response(mock_pipeline)
|
||||
result = pipeline._message_end_to_stream_response()
|
||||
|
||||
assert result.files is not None
|
||||
assert len(result.files) == 1
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from core.repositories.human_input_repository import HumanInputFormSubmissionRepository
|
||||
@@ -68,9 +66,9 @@ def test_pause_reason_payload_carries_approval_channels_through_factory():
|
||||
assert payload.form_token is None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("sqlite_session_factory")
|
||||
def test_enrich_graph_pause_reasons_raises_when_hitl_form_record_is_missing():
|
||||
form_repository = Mock(spec=HumanInputFormSubmissionRepository)
|
||||
form_repository.get_by_form_id.return_value = None
|
||||
form_repository = HumanInputFormSubmissionRepository()
|
||||
|
||||
with pytest.raises(LookupError, match="form-123"):
|
||||
enrich_graph_pause_reasons(
|
||||
|
||||
Reference in new issue
Block a user