diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py index 834cefee87e..b22d5a71148 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py @@ -4,8 +4,8 @@ from __future__ import annotations import json +from dataclasses import dataclass from types import SimpleNamespace -from unittest.mock import Mock import pytest from dify_agent.layers.config import DifyConfigLayerConfig, DifyConfigSkillConfig @@ -32,13 +32,21 @@ from core.app.apps.agent_app.runtime_request_builder import ( AgentAppRuntimeRequestBuilder, AgentAppRuntimeRequestBuildError, ) -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom from core.workflow.file_reference import build_file_reference from graphon.file import File, FileTransferMethod, FileType from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent from graphon.model_runtime.entities.model_entities import ModelFeature from models.agent_config_entities import AgentSoulConfig from tests.unit_tests.config_override import apply_config_overrides +from tests.unit_tests.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema + + +@dataclass(frozen=True) +class _LayerDeps: + """The only part of the agent runtime Deps the user-prompt capability reads.""" + + layers: dict[str, dict[str, object]] @pytest.fixture(autouse=True) @@ -178,7 +186,7 @@ def _ctx( files: tuple[File, ...] = (), image_detail_config: ImagePromptMessageContent.DETAIL | None = None, ) -> AgentAppRuntimeBuildContext: - dify_context = SimpleNamespace( + dify_context = DifyRunContext( tenant_id="tenant-1", app_id="app-1", user_id="user-1", @@ -187,7 +195,7 @@ def _ctx( trace_session_id="trace-session-1", ) return AgentAppRuntimeBuildContext( - dify_context=dify_context, # type: ignore[arg-type] + dify_context=dify_context, agent_id="agent-1", agent_config_snapshot_id="snap-1", agent_soul=soul, @@ -363,10 +371,7 @@ class TestAgentAppRuntimeRequestBuilder: "core.app.apps.agent_app.runtime_request_builder.resolve_model_supports_vision", lambda **_kwargs: vision, ) - model_factory = Mock() - model_factory.return_value.init_model_instance.return_value.get_model_schema.return_value = SimpleNamespace( - features=[ModelFeature.VISION] if vision else [] - ) + model_factory = RecordingModelFactory(schema=llm_schema(features=[ModelFeature.VISION] if vision else [])) monkeypatch.setattr("core.app.knowledge_runtime_config.DifyModelFactory", model_factory) def image_content(file: File, **_kwargs: object) -> ImagePromptMessageContent: @@ -421,7 +426,7 @@ class TestAgentAppRuntimeRequestBuilder: assert "knowledge_fs" in layers config = DifyUserPromptLayerConfig.model_validate(layers["agent_app_user_prompt"].config) # Render the serialized request through the actual runtime module, not just the API DTO. - deps = SimpleNamespace( + deps = _LayerDeps( layers={"agent_app_user_prompt": {"config": json.loads(config.model_dump_json()), "state": {}}} ) prompts = user_prompt_layer.Capability("agent_app_user_prompt").build_user_content(deps) # type: ignore[arg-type] diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py index 37be796bda7..8a3b20630f9 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py @@ -11,6 +11,7 @@ from flask import Flask, current_app from sqlalchemy import Engine, event, select from sqlalchemy.orm import Session, sessionmaker +from core.app.app_config.entities import AppAdditionalFeatures from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueRetrieverResourcesEvent from core.app.entities.task_entities import MessageStreamResponse, StreamEvent, TaskStateMetadata, WorkflowTaskState from core.app.task_pipeline import message_cycle_manager as message_cycle_manager_module @@ -25,6 +26,11 @@ from models.model import App, Conversation, MessageAnnotation, MessageFile from tests.unit_tests.model_factories import make_app, make_conversation +@dataclass(frozen=True) +class _AppConfigWithFeatures: + additional_features: AppAdditionalFeatures + + @dataclass(frozen=True) class _SQLiteDb: engine: Engine @@ -547,10 +553,10 @@ class TestMessageCycleManagerOptimization: def test_knowledge_receipts_survive_message_metadata_without_collapsing_same_document(self, message_cycle_manager): from dify_agent.protocol.knowledge_fs import KnowledgeFsCitation - message_cycle_manager._application_generate_entity.app_config = SimpleNamespace( - additional_features=SimpleNamespace(show_retrieve_source=True) + message_cycle_manager._application_generate_entity.app_config = _AppConfigWithFeatures( + additional_features=AppAdditionalFeatures(show_retrieve_source=True) ) - message_cycle_manager._task_state = SimpleNamespace(metadata=TaskStateMetadata()) + message_cycle_manager._task_state = WorkflowTaskState(metadata=TaskStateMetadata()) sources = [ RetrievalSourceMetadata( dataset_id="space", diff --git a/api/tests/unit_tests/core/app/test_knowledge_runtime_config.py b/api/tests/unit_tests/core/app/test_knowledge_runtime_config.py index 3bc6e51e992..1561a7dfc83 100644 --- a/api/tests/unit_tests/core/app/test_knowledge_runtime_config.py +++ b/api/tests/unit_tests/core/app/test_knowledge_runtime_config.py @@ -1,14 +1,12 @@ -from types import SimpleNamespace -from unittest.mock import Mock - import pytest from dify_agent.protocol.knowledge_fs import KnowledgeFsError from core.app import knowledge_runtime_config as mapping -from core.app.entities.app_invoke_entities import DifyRunContext +from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom from graphon.model_runtime.entities.model_entities import ModelFeature from models.agent_config_entities import AgentSoulConfig from tests.unit_tests.config_override import apply_config_overrides +from tests.unit_tests.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema def soul() -> AgentSoulConfig: @@ -22,23 +20,31 @@ def soul() -> AgentSoulConfig: ) +def run_context() -> DifyRunContext: + return DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ) + + @pytest.mark.parametrize("vision", [True, False]) def test_knowledge_vision_is_owned_by_agent_model_without_provider_env( monkeypatch: pytest.MonkeyPatch, vision: bool ) -> None: apply_config_overrides(monkeypatch, AGENT_SHELL_ENABLED=True) - factory = Mock() - factory.return_value.init_model_instance.return_value.get_model_schema.return_value = SimpleNamespace( - features=[ModelFeature.VISION] if vision else [] - ) + factory = RecordingModelFactory(schema=llm_schema(features=[ModelFeature.VISION] if vision else [])) monkeypatch.setattr(mapping, "DifyModelFactory", factory) - context = Mock(spec=DifyRunContext) + context = run_context() config = mapping.build_knowledge_fs_layer_config(soul(), run_context=context) assert config is not None assert config.agent_supports_vision is vision assert config.spaces[0].name == "Docs" - factory.assert_called_once_with(run_context=context) - factory.return_value.init_model_instance.assert_called_once_with("provider", "model") + assert len(factory.run_contexts) == 1 + assert factory.run_contexts[0] is context + assert factory.model_lookups == [("provider", "model")] assert "provider" not in config.model_dump_json() @@ -46,10 +52,10 @@ def test_no_knowledge_has_no_model_or_shell_dependency_and_import_requires_rebin monkeypatch: pytest.MonkeyPatch, ) -> None: apply_config_overrides(monkeypatch, AGENT_SHELL_ENABLED=False) - assert mapping.build_knowledge_fs_layer_config(AgentSoulConfig(), run_context=Mock(spec=DifyRunContext)) is None + assert mapping.build_knowledge_fs_layer_config(AgentSoulConfig(), run_context=run_context()) is None with pytest.raises(KnowledgeFsError, match="sandbox shell"): - mapping.build_knowledge_fs_layer_config(soul(), run_context=Mock(spec=DifyRunContext)) + mapping.build_knowledge_fs_layer_config(soul(), run_context=run_context()) imported = soul() imported.knowledge.spaces[0].is_missing = True with pytest.raises(KnowledgeFsError, match="Reselect"): - mapping.build_knowledge_fs_layer_config(imported, run_context=Mock(spec=DifyRunContext)) + mapping.build_knowledge_fs_layer_config(imported, run_context=run_context()) diff --git a/api/tests/unit_tests/core/plugin/test_backwards_invocation_datasource.py b/api/tests/unit_tests/core/plugin/test_backwards_invocation_datasource.py index 6af723fbc98..7f09d69f7b1 100644 --- a/api/tests/unit_tests/core/plugin/test_backwards_invocation_datasource.py +++ b/api/tests/unit_tests/core/plugin/test_backwards_invocation_datasource.py @@ -1,5 +1,7 @@ import json -from unittest.mock import MagicMock, patch +from collections.abc import Iterator +from dataclasses import dataclass, field +from unittest.mock import patch import pytest from pydantic import ValidationError @@ -16,6 +18,7 @@ from core.datasource.entities.datasource_entities import ( from core.plugin.backwards_invocation.datasource import PluginDatasourceBackwardsInvocation from core.plugin.entities.request import RequestInvokeDatasource from models.account import Tenant +from tests.unit_tests.knowledge_fs_core_fakes import Marker def _tenant() -> Tenant: @@ -24,10 +27,67 @@ def _tenant() -> Tenant: return tenant -def _tenant() -> Tenant: - tenant = Tenant(name="Test workspace") - tenant.id = "tenant-1" - return tenant +@dataclass +class _RuntimeState: + credentials: dict[str, object] = field(default_factory=dict) + + +@dataclass +class _FakeDatasource: + """Bound datasource plugin double that records each operation call and yields canned messages.""" + + provider_type: DatasourceProviderType + messages: list[object] = field(default_factory=list) + runtime: _RuntimeState = field(default_factory=_RuntimeState) + calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) + + def datasource_provider_type(self) -> DatasourceProviderType: + return self.provider_type + + def _record(self, operation: str, kwargs: dict[str, object]) -> Iterator[object]: + self.calls.append((operation, kwargs)) + return iter(self.messages) + + def get_website_crawl(self, **kwargs: object) -> Iterator[object]: + return self._record("get_website_crawl", kwargs) + + def get_online_document_pages(self, **kwargs: object) -> Iterator[object]: + return self._record("get_online_document_pages", kwargs) + + def get_online_document_page_content(self, **kwargs: object) -> Iterator[object]: + return self._record("get_online_document_page_content", kwargs) + + def online_drive_browse_files(self, **kwargs: object) -> Iterator[object]: + return self._record("online_drive_browse_files", kwargs) + + def online_drive_download_file(self, **kwargs: object) -> Iterator[object]: + return self._record("online_drive_download_file", kwargs) + + +@dataclass(frozen=True) +class _ProviderEntity: + provider_type: DatasourceProviderType + + +@dataclass +class _FakeController: + entity: _ProviderEntity + datasource: _FakeDatasource | None = None + need_credentials: bool = True + requested_datasources: list[str] = field(default_factory=list) + validations: list[dict[str, object]] = field(default_factory=list) + + def get_datasource(self, datasource_name: str) -> _FakeDatasource: + self.requested_datasources.append(datasource_name) + assert self.datasource is not None + return self.datasource + + def validate_credentials(self, **kwargs: object) -> None: + self.validations.append(kwargs) + + +def _controller(datasource: _FakeDatasource) -> _FakeController: + return _FakeController(entity=_ProviderEntity(datasource.provider_type), datasource=datasource) def datasource_payload(**overrides: object) -> RequestInvokeDatasource: @@ -58,14 +118,8 @@ def test_request_rejects_operation_and_datasource_type_mismatch() -> None: def test_invoke_online_document_uses_dify_bound_runtime_and_credential() -> None: payload = datasource_payload() message = OnlineDocumentPagesMessage(result=[]) - runtime = MagicMock() - runtime.runtime.credentials = dict[str, object]() - runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DOCUMENT - runtime.get_online_document_pages.return_value = iter([message]) - controller = MagicMock() - controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT - controller.need_credentials = True - controller.get_datasource.return_value = runtime + runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT, messages=[message]) + controller = _controller(runtime) with ( patch( @@ -98,11 +152,17 @@ def test_invoke_online_document_uses_dify_bound_runtime_and_credential() -> None credential_id="credential-1", ) assert runtime.runtime.credentials == {"token": "resolved-by-dify"} - runtime.get_online_document_pages.assert_called_once_with( - user_id="user-1", - datasource_parameters={"cursor": "next"}, - provider_type=DatasourceProviderType.ONLINE_DOCUMENT, - ) + assert controller.requested_datasources == ["notion_datasource"] + assert runtime.calls == [ + ( + "get_online_document_pages", + { + "user_id": "user-1", + "datasource_parameters": {"cursor": "next"}, + "provider_type": DatasourceProviderType.ONLINE_DOCUMENT, + }, + ) + ] def test_validate_credentials_uses_resolved_provider_controller() -> None: @@ -110,11 +170,8 @@ def test_validate_credentials_uses_resolved_provider_controller() -> None: datasource_parameters={}, operation="validate_credentials", ) - runtime = MagicMock() - controller = MagicMock() - controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT - controller.need_credentials = True - controller.get_datasource.return_value = runtime + runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT) + controller = _controller(runtime) with ( patch( @@ -135,10 +192,8 @@ def test_validate_credentials_uses_resolved_provider_controller() -> None: ) assert result == [{"result": True}] - controller.validate_credentials.assert_called_once_with( - user_id="user-1", - credentials={"token": "resolved-by-dify"}, - ) + assert controller.validations == [{"user_id": "user-1", "credentials": {"token": "resolved-by-dify"}}] + assert runtime.calls == [] def test_invoke_website_crawl_uses_the_bound_datasource_plugin() -> None: @@ -149,15 +204,9 @@ def test_invoke_website_crawl_uses_the_bound_datasource_plugin() -> None: operation="get_website_crawl", provider="langgenius/firecrawl_datasource/firecrawl", ) - message = MagicMock() - runtime = MagicMock() - runtime.runtime.credentials = dict[str, object]() - runtime.datasource_provider_type.return_value = DatasourceProviderType.WEBSITE_CRAWL - runtime.get_website_crawl.return_value = iter([message]) - controller = MagicMock() - controller.entity.provider_type = DatasourceProviderType.WEBSITE_CRAWL - controller.need_credentials = True - controller.get_datasource.return_value = runtime + message = Marker("datasource-message") + runtime = _FakeDatasource(DatasourceProviderType.WEBSITE_CRAWL, messages=[message]) + controller = _controller(runtime) with ( patch( @@ -179,11 +228,16 @@ def test_invoke_website_crawl_uses_the_bound_datasource_plugin() -> None: assert result == [message] assert runtime.runtime.credentials == {"api_key": "resolved-by-dify"} - runtime.get_website_crawl.assert_called_once_with( - user_id="user-1", - datasource_parameters={"url": "https://example.com"}, - provider_type=DatasourceProviderType.WEBSITE_CRAWL, - ) + assert runtime.calls == [ + ( + "get_website_crawl", + { + "user_id": "user-1", + "datasource_parameters": {"url": "https://example.com"}, + "provider_type": DatasourceProviderType.WEBSITE_CRAWL, + }, + ) + ] def test_invoke_online_document_content_builds_the_dify_request_entity() -> None: @@ -193,15 +247,9 @@ def test_invoke_online_document_content_builds_the_dify_request_entity() -> None operation="get_online_document_page_content", page=page, ) - message = MagicMock() - runtime = MagicMock() - runtime.runtime.credentials = dict[str, object]() - runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DOCUMENT - runtime.get_online_document_page_content.return_value = iter([message]) - controller = MagicMock() - controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT - controller.need_credentials = True - controller.get_datasource.return_value = runtime + message = Marker("datasource-message") + runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT, messages=[message]) + controller = _controller(runtime) with ( patch( @@ -222,11 +270,16 @@ def test_invoke_online_document_content_builds_the_dify_request_entity() -> None ) assert result == [message] - runtime.get_online_document_page_content.assert_called_once_with( - user_id="user-1", - datasource_parameters=page, - provider_type=DatasourceProviderType.ONLINE_DOCUMENT, - ) + assert runtime.calls == [ + ( + "get_online_document_page_content", + { + "user_id": "user-1", + "datasource_parameters": page, + "provider_type": DatasourceProviderType.ONLINE_DOCUMENT, + }, + ) + ] @pytest.mark.parametrize( @@ -260,18 +313,9 @@ def test_invoke_online_drive_builds_typed_dify_requests( provider="langgenius/google_drive/google_drive", request=request_data, ) - message = MagicMock() - runtime = MagicMock() - runtime.runtime.credentials = dict[str, object]() - runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DRIVE - { - "online_drive_browse_files": runtime.online_drive_browse_files, - "online_drive_download_file": runtime.online_drive_download_file, - }[method_name].return_value = iter([message]) - controller = MagicMock() - controller.entity.provider_type = DatasourceProviderType.ONLINE_DRIVE - controller.need_credentials = True - controller.get_datasource.return_value = runtime + message = Marker("datasource-message") + runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DRIVE, messages=[message]) + controller = _controller(runtime) with ( patch( @@ -292,13 +336,12 @@ def test_invoke_online_drive_builds_typed_dify_requests( ) assert result == [message] - call = { - "online_drive_browse_files": runtime.online_drive_browse_files, - "online_drive_download_file": runtime.online_drive_download_file, - }[method_name].call_args - assert call.kwargs["user_id"] == "user-1" - assert call.kwargs["provider_type"] == DatasourceProviderType.ONLINE_DRIVE - assert call.kwargs["request"] == request_type.model_validate(request_data) + assert len(runtime.calls) == 1 + called_operation, call_kwargs = runtime.calls[0] + assert called_operation == method_name + assert call_kwargs["user_id"] == "user-1" + assert call_kwargs["provider_type"] == DatasourceProviderType.ONLINE_DRIVE + assert call_kwargs["request"] == request_type.model_validate(request_data) def test_online_drive_metadata_and_binary_survive_the_inner_event_stream() -> None: @@ -397,10 +440,7 @@ def test_online_drive_metadata_and_binary_survive_the_inner_event_stream() -> No def test_invoke_rejects_missing_required_credential() -> None: payload = datasource_payload() - controller = MagicMock() - controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT - controller.need_credentials = True - controller.get_datasource.return_value = MagicMock() + controller = _controller(_FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT)) with ( patch( @@ -424,8 +464,7 @@ def test_invoke_rejects_missing_required_credential() -> None: def test_invoke_rejects_provider_declaration_type_mismatch() -> None: payload = datasource_payload() - controller = MagicMock() - controller.entity.provider_type = DatasourceProviderType.ONLINE_DRIVE + controller = _FakeController(entity=_ProviderEntity(DatasourceProviderType.ONLINE_DRIVE)) with ( patch( diff --git a/api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py b/api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py index a5142ec8d1b..8053e25c64a 100644 --- a/api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py +++ b/api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py @@ -1,6 +1,6 @@ import json -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from dataclasses import dataclass, field +from unittest.mock import patch import pytest @@ -129,24 +129,111 @@ def test_invoke_summary_uses_same_user_scope_for_token_helpers(): ) +def _tenant() -> Tenant: + tenant = Tenant(name="Test Workspace") + tenant.id = "tenant-1" + return tenant + + +@dataclass +class _RecordingModelInstance: + """Bound model instance double: returns a canned result and records invocation kwargs.""" + + result: str + calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) + + def invoke_text_embedding(self, **kwargs: object) -> str: + self.calls.append(("invoke_text_embedding", kwargs)) + return self.result + + def invoke_multimodal_embedding(self, **kwargs: object) -> str: + self.calls.append(("invoke_multimodal_embedding", kwargs)) + return self.result + + def invoke_rerank(self, **kwargs: object) -> str: + self.calls.append(("invoke_rerank", kwargs)) + return self.result + + +@dataclass(frozen=True) +class _ModelProvider: + provider: str + + +@dataclass(frozen=True) +class _ActiveModel: + deprecated: bool + features: list[ModelFeature] + fetch_from: str + label: dict[str, str] + model: str + model_properties: dict[str, int] + model_type: ModelType + provider: _ModelProvider + status: str + + +@dataclass +class _ProviderConfigurations: + models: list[_ActiveModel] + get_models_calls: list[dict[str, object]] = field(default_factory=list) + + def get_models(self, **kwargs: object) -> list[_ActiveModel]: + self.get_models_calls.append(kwargs) + return self.models + + +@dataclass(frozen=True) +class _ProviderManager: + configurations: _ProviderConfigurations + + def get_configurations(self, _tenant_id: str) -> _ProviderConfigurations: + return self.configurations + + +@dataclass(frozen=True) +class _InstalledPlugin: + plugin_id: str + plugin_unique_identifier: str + + +@dataclass(frozen=True) +class _ParameterRule: + name: str + use_template: str | None + max: object + default: object = None + + +@dataclass(frozen=True) +class _ModelSchema: + model_properties: dict[str, int] + parameter_rules: list[_ParameterRule] + + +@dataclass(frozen=True) +class _SchemaModelInstance: + schema: _ModelSchema + + def get_model_schema(self) -> _ModelSchema: + return self.schema + + def test_invoke_text_embedding_forwards_input_type_to_bound_model_instance(): - tenant = SimpleNamespace(id="tenant-1") + tenant = _tenant() payload = RequestInvokeTextEmbedding( provider="langgenius/openai/openai", model="text-embedding-3-small", texts=["query text"], input_type=EmbeddingInputType.QUERY, ) - model_instance = SimpleNamespace(invoke_text_embedding=lambda **_: "unused") + model_instance = _RecordingModelInstance(result="embedding-result") - with ( - patch.object( - PluginModelBackwardsInvocation, - "_get_bound_model_instance", - return_value=model_instance, - ) as mock_get_model, - patch.object(model_instance, "invoke_text_embedding", return_value="embedding-result") as mock_invoke, - ): + with patch.object( + PluginModelBackwardsInvocation, + "_get_bound_model_instance", + return_value=model_instance, + ) as mock_get_model: result = PluginModelBackwardsInvocation.invoke_text_embedding("user-1", tenant, payload) assert result == "embedding-result" @@ -158,11 +245,13 @@ def test_invoke_text_embedding_forwards_input_type_to_bound_model_instance(): model="text-embedding-3-small", request_metadata={}, ) - mock_invoke.assert_called_once_with(texts=["query text"], input_type=EmbeddingInputType.QUERY) + assert model_instance.calls == [ + ("invoke_text_embedding", {"texts": ["query text"], "input_type": EmbeddingInputType.QUERY}) + ] def test_invoke_multimodal_embedding_uses_bound_model_instance(): - tenant = SimpleNamespace(id="tenant-1") + tenant = _tenant() documents = [{"content": "AQID", "content_type": "image", "file_id": "asset.png"}] payload = RequestInvokeMultimodalEmbedding( provider="langgenius/clip/clip", @@ -170,20 +259,13 @@ def test_invoke_multimodal_embedding_uses_bound_model_instance(): documents=documents, input_type=EmbeddingInputType.DOCUMENT, ) - model_instance = SimpleNamespace(invoke_multimodal_embedding=lambda **_: "unused") + model_instance = _RecordingModelInstance(result="embedding-result") - with ( - patch.object( - PluginModelBackwardsInvocation, - "_get_bound_model_instance", - return_value=model_instance, - ) as mock_get_model, - patch.object( - model_instance, - "invoke_multimodal_embedding", - return_value="embedding-result", - ) as mock_invoke, - ): + with patch.object( + PluginModelBackwardsInvocation, + "_get_bound_model_instance", + return_value=model_instance, + ) as mock_get_model: result = PluginModelBackwardsInvocation.invoke_multimodal_embedding("user-1", tenant, payload) assert result == "embedding-result" @@ -195,39 +277,38 @@ def test_invoke_multimodal_embedding_uses_bound_model_instance(): model="clip-multimodal", request_metadata={}, ) - mock_invoke.assert_called_once_with( - multimodel_documents=documents, - input_type=EmbeddingInputType.DOCUMENT, - ) + assert model_instance.calls == [ + ( + "invoke_multimodal_embedding", + {"multimodel_documents": documents, "input_type": EmbeddingInputType.DOCUMENT}, + ) + ] def test_invoke_rerank_accepts_omitted_threshold_and_top_n(): - tenant = SimpleNamespace(id="tenant-1") + tenant = _tenant() payload = RequestInvokeRerank( provider="langgenius/cohere/cohere", model="rerank-v3.5", query="query", docs=["first", "second"], ) - model_instance = SimpleNamespace(invoke_rerank=lambda **_: "unused") + model_instance = _RecordingModelInstance(result="rerank-result") - with ( - patch.object( - PluginModelBackwardsInvocation, - "_get_bound_model_instance", - return_value=model_instance, - ), - patch.object(model_instance, "invoke_rerank", return_value="rerank-result") as mock_invoke, + with patch.object( + PluginModelBackwardsInvocation, + "_get_bound_model_instance", + return_value=model_instance, ): result = PluginModelBackwardsInvocation.invoke_rerank("user-1", tenant, payload) assert result == "rerank-result" - mock_invoke.assert_called_once_with( - query="query", - docs=["first", "second"], - score_threshold=None, - top_n=None, - ) + assert model_instance.calls == [ + ( + "invoke_rerank", + {"query": "query", "docs": ["first", "second"], "score_threshold": None, "top_n": None}, + ) + ] @pytest.mark.parametrize( @@ -237,8 +318,8 @@ def test_invoke_rerank_accepts_omitted_threshold_and_top_n(): def test_list_models_returns_only_active_tenant_models_with_installed_identity(model_type, selected): model_name = "selected-model" payload = RequestListModels(model_type=model_type, model=model_name if selected else None, limit=10) - provider = SimpleNamespace(provider="langgenius/openai/openai") - active_model = SimpleNamespace( + provider = _ModelProvider(provider="langgenius/openai/openai") + active_model = _ActiveModel( deprecated=False, features=[ModelFeature.VISION], fetch_from="predefined-model", @@ -249,9 +330,9 @@ def test_list_models_returns_only_active_tenant_models_with_installed_identity(m provider=provider, status="active", ) - configurations = SimpleNamespace(get_models=MagicMock(return_value=[active_model])) - provider_manager = SimpleNamespace(get_configurations=lambda _tenant_id: configurations) - installed_plugin = SimpleNamespace( + configurations = _ProviderConfigurations(models=[active_model]) + provider_manager = _ProviderManager(configurations=configurations) + installed_plugin = _InstalledPlugin( plugin_id="langgenius/openai", plugin_unique_identifier="langgenius/openai:1.0.0@sha256:installed", ) @@ -278,10 +359,7 @@ def test_list_models_returns_only_active_tenant_models_with_installed_identity(m assert result.items[0].provider == "openai" assert result.items[0].plugin_unique_identifier == installed_plugin.plugin_unique_identifier assert result.items[0].capabilities["features"] == ["vision"] - configurations.get_models.assert_called_once_with( - model_type=model_type, - only_active=True, - ) + assert configurations.get_models_calls == [{"model_type": model_type, "only_active": True}] if selected: resolve_limits.assert_called_once_with( tenant_id="tenant-1", user_id="user-1", provider=provider.provider, model=model_name @@ -298,11 +376,11 @@ def test_list_models_returns_only_active_tenant_models_with_installed_identity(m @pytest.mark.parametrize("parameter", ["max_tokens", "max_completion_tokens", "maxtoken", "custom_limit"]) def test_model_token_limits_use_schema_maximum_not_default(parameter): - schema = SimpleNamespace( + schema = _ModelSchema( model_properties={"context_size": 131072}, - parameter_rules=[SimpleNamespace(name=parameter, use_template="max_tokens", max=32768.0, default=1024)], + parameter_rules=[_ParameterRule(name=parameter, use_template="max_tokens", max=32768.0, default=1024)], ) - instance = SimpleNamespace(get_model_schema=lambda: schema) + instance = _SchemaModelInstance(schema=schema) with patch.object(PluginModelBackwardsInvocation, "_get_bound_model_instance", return_value=instance) as bound: limits = PluginModelBackwardsInvocation._model_token_limits( tenant_id="tenant-1", user_id="user-1", provider="plugin/provider", model="selected-model" @@ -321,14 +399,14 @@ def test_model_token_limits_use_schema_maximum_not_default(parameter): @pytest.mark.parametrize("maximum", [None, 0, -1, True, float("inf"), float("nan"), 10**400, 12.5, "32768"]) def test_model_token_limits_do_not_invent_unknown_output_capacity(maximum): - schema = SimpleNamespace( + schema = _ModelSchema( model_properties={"context_size": 8192}, - parameter_rules=[SimpleNamespace(name="max_tokens", use_template=None, max=maximum, default=4096)], + parameter_rules=[_ParameterRule(name="max_tokens", use_template=None, max=maximum, default=4096)], ) with patch.object( PluginModelBackwardsInvocation, "_get_bound_model_instance", - return_value=SimpleNamespace(get_model_schema=lambda: schema), + return_value=_SchemaModelInstance(schema=schema), ): limits = PluginModelBackwardsInvocation._model_token_limits( tenant_id="tenant-1", user_id="user-1", provider="plugin/provider", model="selected-model" @@ -338,17 +416,17 @@ def test_model_token_limits_do_not_invent_unknown_output_capacity(maximum): def test_model_token_limits_respect_the_tighter_declared_alias(): - schema = SimpleNamespace( + schema = _ModelSchema( model_properties={}, parameter_rules=[ - SimpleNamespace(name="max_tokens", use_template=None, max=32768), - SimpleNamespace(name="max_output_tokens", use_template=None, max=8192), + _ParameterRule(name="max_tokens", use_template=None, max=32768), + _ParameterRule(name="max_output_tokens", use_template=None, max=8192), ], ) with patch.object( PluginModelBackwardsInvocation, "_get_bound_model_instance", - return_value=SimpleNamespace(get_model_schema=lambda: schema), + return_value=_SchemaModelInstance(schema=schema), ): limits = PluginModelBackwardsInvocation._model_token_limits( tenant_id="tenant-1", user_id="user-1", provider="plugin/provider", model="selected-model" diff --git a/api/tests/unit_tests/core/plugin/test_model_billing_attribution.py b/api/tests/unit_tests/core/plugin/test_model_billing_attribution.py index 2accc76c71f..ed78a80da45 100644 --- a/api/tests/unit_tests/core/plugin/test_model_billing_attribution.py +++ b/api/tests/unit_tests/core/plugin/test_model_billing_attribution.py @@ -1,13 +1,20 @@ """Exercise backwards invocation through the real hosted quota owner, with no external I/O.""" from collections.abc import Callable -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from dataclasses import dataclass +from unittest.mock import patch import pytest from configs import dify_config -from core.entities.provider_entities import ProviderQuotaType, QuotaUnit +from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle +from core.entities.provider_entities import ( + CustomConfiguration, + ProviderQuotaType, + QuotaConfiguration, + QuotaUnit, + SystemConfiguration, +) from core.model_context import get_credit_usage_metadata from core.model_manager import ModelInstance from core.plugin.backwards_invocation.model import PluginModelBackwardsInvocation @@ -17,12 +24,80 @@ from core.plugin.entities.request import ( RequestInvokeRerank, RequestInvokeTextEmbedding, ) +from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime from enums import DeploymentEdition +from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage from graphon.model_runtime.entities.model_entities import ModelType +from graphon.model_runtime.entities.provider_entities import ProviderEntity +from graphon.model_runtime.model_providers.base.ai_model import AIModel +from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from graphon.model_runtime.model_providers.base.rerank_model import RerankModel +from graphon.model_runtime.model_providers.base.text_embedding_model import TextEmbeddingModel +from models.account import Tenant from models.provider import ProviderType from services.billing_service import BillingService +from tests.unit_tests.knowledge_fs_core_fakes import Marker + +_MODEL_TYPE_CLASSES: dict[ModelType, type[AIModel]] = { + ModelType.LLM: LargeLanguageModel, + ModelType.TEXT_EMBEDDING: TextEmbeddingModel, + ModelType.RERANK: RerankModel, +} + + +@dataclass(frozen=True) +class _ProviderManager: + bundle: ProviderModelBundle + + def get_provider_model_bundle(self, **_kwargs: object) -> ProviderModelBundle: + return self.bundle + + def get_configurations(self, _tenant_id: str) -> dict[str, ProviderConfiguration]: + return {self.bundle.configuration.provider.provider: self.bundle.configuration} + + +def _tenant() -> Tenant: + tenant = Tenant(name="Test Workspace") + tenant.id = "tenant-1" + return tenant + + +def _hosted_bundle(model_type: ModelType) -> ProviderModelBundle: + provider = ProviderEntity( + provider="langgenius/openai/openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[model_type], + configurate_methods=[], + ) + configuration = ProviderConfiguration( + tenant_id="tenant-1", + provider=provider, + preferred_provider_type=ProviderType.SYSTEM, + using_provider_type=ProviderType.SYSTEM, + system_configuration=SystemConfiguration( + enabled=True, + current_quota_type=ProviderQuotaType.TRIAL, + quota_configurations=[ + QuotaConfiguration( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.CREDITS, + quota_limit=200, + quota_used=0, + is_valid=True, + restrict_models=[], + ) + ], + ), + custom_configuration=CustomConfiguration(provider=None), + model_settings=[], + ) + model_type_instance = _MODEL_TYPE_CLASSES[model_type]( + provider_schema=provider, + model_runtime=create_plugin_model_runtime(tenant_id="tenant-1", user_id="synthetic-end-user"), + ) + return ProviderModelBundle(configuration=configuration, model_type_instance=model_type_instance) @pytest.mark.parametrize( @@ -73,29 +148,7 @@ def test_model_invocation_preserves_billing_origin( config_overrides: Callable[..., None], ) -> None: config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - configuration = SimpleNamespace( - tenant_id="tenant-1", - provider=SimpleNamespace(provider="langgenius/openai/openai"), - model_settings=[], - using_provider_type=ProviderType.SYSTEM, - system_configuration=SimpleNamespace( - current_quota_type=ProviderQuotaType.TRIAL, - quota_configurations=[ - SimpleNamespace( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.CREDITS, - quota_limit=200, - restrict_models=[], - ) - ], - ), - get_current_credentials=lambda **_: {"api_key": "local-test-only"}, - get_provider_model=lambda **_: None, - ) - bundle = SimpleNamespace(configuration=configuration, model_type_instance=MagicMock(model_type=model_type)) - provider_manager = MagicMock() - provider_manager.get_provider_model_bundle.return_value = bundle - provider_manager.get_configurations.return_value = {"langgenius/openai/openai": configuration} + provider_manager = _ProviderManager(bundle=_hosted_bundle(model_type)) metadata = {"app_type": "new_rag", "created_by": created_by} if created_by else {} payload = payload_type.model_validate( {"provider": "langgenius/openai/openai", "model": "selected-model", **body, **metadata} @@ -117,15 +170,17 @@ def test_model_invocation_preserves_billing_origin( index=0, message=AssistantPromptMessage(content="OK"), usage=LLMUsage.empty_usage(), finish_reason="stop" ), ) - provider_result = (item for item in [chunk]) if model_type == ModelType.LLM else MagicMock() + provider_result = (item for item in [chunk]) if model_type == ModelType.LLM else Marker("provider-result") with ( patch("core.model_manager.create_plugin_provider_manager", return_value=provider_manager), patch("core.app.llm.quota.create_plugin_provider_manager", return_value=provider_manager), patch.object(type(dify_config), "get_model_credits", return_value=3), + patch.object(ProviderConfiguration, "get_current_credentials", return_value={"api_key": "local-test-only"}), + patch.object(ProviderConfiguration, "get_provider_model", return_value=None), patch.object(ModelInstance, provider_method, return_value=provider_result), patch.object(BillingService, "_send_quota_request", side_effect=billing_transport), ): - response = invoke_method("synthetic-end-user", SimpleNamespace(id="tenant-1"), payload) + response = invoke_method("synthetic-end-user", _tenant(), payload) assert get_credit_usage_metadata() is None if model_type == ModelType.LLM: assert list(response)[0].delta.message.content == "OK" diff --git a/api/tests/unit_tests/core/tools/builtin_tool/providers/knowledge_fs/test_create_research_task.py b/api/tests/unit_tests/core/tools/builtin_tool/providers/knowledge_fs/test_create_research_task.py index 2ea44119a14..09813767df0 100644 --- a/api/tests/unit_tests/core/tools/builtin_tool/providers/knowledge_fs/test_create_research_task.py +++ b/api/tests/unit_tests/core/tools/builtin_tool/providers/knowledge_fs/test_create_research_task.py @@ -1,6 +1,6 @@ from __future__ import annotations -from types import SimpleNamespace +from dataclasses import dataclass from unittest.mock import patch import pytest @@ -41,6 +41,18 @@ class AppCapabilities: ) +class RevokedAppCapabilities: + def create_research_task(self, **_kwargs: object) -> KnowledgeFSResearchTaskResponse: + raise ToolInvokeError("binding revoked") + + +@dataclass(frozen=True) +class KnowledgeFSRuntimeStub: + """The only KnowledgeFS runtime component this tool reaches.""" + + app_capabilities: AppCapabilities | RevokedAppCapabilities + + def _run_context() -> DifyRunContext: return DifyRunContext( tenant_id="tenant-from-context", @@ -103,7 +115,7 @@ def test_tool_is_a_real_agent_and_workflow_consumer_using_run_context_identity( caller_kind: KnowledgeFSAppSpaceJoinType, ) -> None: capabilities = AppCapabilities() - runtime = SimpleNamespace(app_capabilities=capabilities) + runtime = KnowledgeFSRuntimeStub(app_capabilities=capabilities) tool = _tool(tool_invoke_from=tool_invoke_from, run_context=_run_context()) engine = create_engine("sqlite://") @@ -212,7 +224,7 @@ def test_tool_rejects_cross_tenant_and_non_application_callers_before_runtime_io def test_tool_omits_unspecified_mode_and_maps_invalid_model_input_to_tool_error() -> None: capabilities = AppCapabilities() - runtime = SimpleNamespace(app_capabilities=capabilities) + runtime = KnowledgeFSRuntimeStub(app_capabilities=capabilities) tool = _tool(tool_invoke_from=ToolInvokeFrom.AGENT, run_context=_run_context()) engine = create_engine("sqlite://") @@ -256,11 +268,7 @@ def test_tool_omits_unspecified_mode_and_maps_invalid_model_input_to_tool_error( def test_tool_preserves_explicit_tool_errors_from_the_capability_boundary() -> None: tool = _tool(tool_invoke_from=ToolInvokeFrom.AGENT, run_context=_run_context()) - runtime = SimpleNamespace( - app_capabilities=SimpleNamespace( - create_research_task=lambda **_: (_ for _ in ()).throw(ToolInvokeError("binding revoked")) - ) - ) + runtime = KnowledgeFSRuntimeStub(app_capabilities=RevokedAppCapabilities()) engine = create_engine("sqlite://") with ( diff --git a/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval_v2/test_automatic_metadata_filter.py b/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval_v2/test_automatic_metadata_filter.py index 3bcf675bffa..a0f3fd54222 100644 --- a/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval_v2/test_automatic_metadata_filter.py +++ b/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval_v2/test_automatic_metadata_filter.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from types import SimpleNamespace +from dataclasses import dataclass, field import pytest @@ -29,6 +29,14 @@ from graphon.nodes.llm.entities import ModelConfig from services.knowledge_fs.product_dto import KnowledgeFSMetadataFieldResponse +@dataclass(frozen=True) +class _ModelConfigEntity: + """The prompt and invocation settings the extractor reads from the resolved model config.""" + + stop: list[str] + parameters: dict[str, object] = field(default_factory=dict) + + def _field(name: str, field_type: str) -> KnowledgeFSMetadataFieldResponse: return KnowledgeFSMetadataFieldResponse.model_validate( { @@ -61,7 +69,7 @@ def test_intersect_metadata_fields_keeps_only_fields_shared_with_identical_types def test_build_metadata_filter_prompt_mirrors_legacy_chat_and_completion_templates() -> None: - model_config = SimpleNamespace(stop=[""]) + model_config = _ModelConfigEntity(stop=[""]) chat_messages, chat_stop = build_metadata_filter_prompt( model_config=model_config, # type: ignore[arg-type] @@ -206,7 +214,7 @@ def test_extractor_invokes_the_configured_model_with_the_legacy_prompt_and_credi '{"metadata_map": [{"metadata_field_name": "department", "metadata_field_value": "finance", ' '"comparison_operator": "="}]}' ) - model_config_entity = SimpleNamespace(parameters={"temperature": 0.2}, stop=["END"]) + model_config_entity = _ModelConfigEntity(parameters={"temperature": 0.2}, stop=["END"]) fetched: list[dict[str, object]] = [] def fake_fetch_model_config(**kwargs: object) -> tuple[object, object]: diff --git a/api/tests/unit_tests/core/workflow/test_node_runtime.py b/api/tests/unit_tests/core/workflow/test_node_runtime.py index 472809744c9..189a1d1208e 100644 --- a/api/tests/unit_tests/core/workflow/test_node_runtime.py +++ b/api/tests/unit_tests/core/workflow/test_node_runtime.py @@ -410,13 +410,24 @@ def test_dify_prepared_polling_llm_commits_successful_reservation() -> None: reservation.release.assert_not_called() -def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_result() -> None: +def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: succeeded_result = LLMPollingResult.model_construct(status=LLMPollingStatus.SUCCEEDED, result=None) - polling_runtime = SimpleNamespace( - start_llm_polling=Mock(return_value=succeeded_result), - check_llm_polling=Mock(), + polling_runtime = PluginModelRuntime( + tenant_id="tenant-id", + user_id="user-id", + client=PluginModelClient(), + plugin_service=PluginService, ) - reservation = MagicMock() + started_pollings: list[dict[str, object]] = [] + + def start_llm_polling(**kwargs: object) -> LLMPollingResult: + started_pollings.append(kwargs) + return succeeded_result + + monkeypatch.setattr(polling_runtime, "start_llm_polling", start_llm_polling) + reservation = sentinel.quota_reservation model_instance = _QuotaManagedModelInstanceStub( model_schema=_build_model_schema(features=[ModelFeature.POLLING]), model_runtime=polling_runtime, @@ -437,6 +448,7 @@ def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_resu ) release_quota_safely.assert_called_once_with(reservation) + assert len(started_pollings) == 1 def test_dify_prepared_polling_llm_releases_previous_reservation_on_restart() -> None: diff --git a/api/tests/unit_tests/knowledge_fs_core_fakes.py b/api/tests/unit_tests/knowledge_fs_core_fakes.py index c6f4aa66eca..cca51b19b0f 100644 --- a/api/tests/unit_tests/knowledge_fs_core_fakes.py +++ b/api/tests/unit_tests/knowledge_fs_core_fakes.py @@ -8,6 +8,8 @@ from __future__ import annotations from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from graphon.model_runtime.entities.common_entities import I18nObject +from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelFeature, ModelType from services.knowledge_fs.observability import ( KnowledgeFSBatchStatusMetric, KnowledgeFSCapabilityIssuanceMetric, @@ -100,3 +102,39 @@ class RecordingOperationalMetrics: def register_control_space_state_gauge(self, read_counts: Callable[[], Mapping[str, int]]) -> None: self.gauge_readers.append(read_counts) + + +def llm_schema(*, features: list[ModelFeature]) -> AIModelEntity: + return AIModelEntity( + model="model", + label=I18nObject(en_US="model"), + model_type=ModelType.LLM, + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={}, + features=features, + ) + + +@dataclass +class SchemaModelInstance: + schema: AIModelEntity + + def get_model_schema(self) -> AIModelEntity: + return self.schema + + +@dataclass +class RecordingModelFactory: + """Stands in for DifyModelFactory and records the constructor context and model lookups.""" + + schema: AIModelEntity + run_contexts: list[object] = field(default_factory=list) + model_lookups: list[tuple[str, str]] = field(default_factory=list) + + def __call__(self, *, run_context: object) -> RecordingModelFactory: + self.run_contexts.append(run_context) + return self + + def init_model_instance(self, provider: str, model: str) -> SchemaModelInstance: + self.model_lookups.append((provider, model)) + return SchemaModelInstance(self.schema) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_cutover_runtime_gate.py b/api/tests/unit_tests/services/test_knowledge_fs_cutover_runtime_gate.py index 9a3f48d1805..aec3b62f379 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_cutover_runtime_gate.py +++ b/api/tests/unit_tests/services/test_knowledge_fs_cutover_runtime_gate.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Callable from datetime import UTC, datetime import pytest @@ -12,7 +13,6 @@ from models.knowledge_fs_cutover import ( ) from services.knowledge_fs.cutover_runtime_gate import SQLKnowledgeFSWorkspaceRuntimeGate from services.knowledge_fs.product_remote import KnowledgeFSOperationUnavailableError -from tests.unit_tests.knowledge_fs_core_fakes import RecordingCallable _WATERMARK: KnowledgeFSCutoverRevisionWatermark = { "membership_epoch": 1, @@ -34,39 +34,43 @@ def test_runtime_gate_fails_closed_without_a_cutover_ledger(sqlite_session: Sess class _Initializer: - def __init__(self, ensure_initialized: RecordingCallable) -> None: - self.ensure_initialized = ensure_initialized + def __init__(self, initialize: Callable[..., KnowledgeFSWorkspaceCutoverLedger]) -> None: + self._initialize = initialize + self.calls: list[str] = [] + + def ensure_initialized(self, *, tenant_id: str) -> KnowledgeFSWorkspaceCutoverLedger: + self.calls.append(tenant_id) + return self._initialize(tenant_id=tenant_id) @pytest.mark.parametrize("sqlite_session", [(KnowledgeFSWorkspaceCutoverLedger,)], indirect=True) def test_runtime_gate_initializes_a_missing_greenfield_workspace_once(sqlite_session: Session) -> None: session_maker = sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) - def initialize(*, tenant_id: str) -> None: + def initialize(*, tenant_id: str) -> KnowledgeFSWorkspaceCutoverLedger: + ledger = KnowledgeFSWorkspaceCutoverLedger( + tenant_id=tenant_id, + source_revision_watermark=_WATERMARK, + applied_revision_watermark=_WATERMARK, + phase=KnowledgeFSWorkspaceCutoverPhase.CUTOVER, + freeze_at=datetime(2026, 7, 21, 11, 59, tzinfo=UTC), + cutover_at=datetime(2026, 7, 21, 12, 0, tzinfo=UTC), + product_routes_enabled=True, + capability_v2_enabled=True, + integrated_mode_enabled=True, + legacy_acl_read_only=True, + ) with session_maker.begin() as session: - session.add( - KnowledgeFSWorkspaceCutoverLedger( - tenant_id=tenant_id, - source_revision_watermark=_WATERMARK, - applied_revision_watermark=_WATERMARK, - phase=KnowledgeFSWorkspaceCutoverPhase.CUTOVER, - freeze_at=datetime(2026, 7, 21, 11, 59, tzinfo=UTC), - cutover_at=datetime(2026, 7, 21, 12, 0, tzinfo=UTC), - product_routes_enabled=True, - capability_v2_enabled=True, - integrated_mode_enabled=True, - legacy_acl_read_only=True, - ) - ) + session.add(ledger) + return ledger - ensure_initialized = RecordingCallable(side_effect=initialize) - initializer = _Initializer(ensure_initialized) + initializer = _Initializer(initialize) gate = SQLKnowledgeFSWorkspaceRuntimeGate(session_maker, initializer=initializer) gate.require_product_routes(tenant_id="tenant-1") gate.require_capability_v2(tenant_id="tenant-1") - ensure_initialized.assert_called_once_with(tenant_id="tenant-1") + assert initializer.calls == ["tenant-1"] @pytest.mark.parametrize("sqlite_session", [(KnowledgeFSWorkspaceCutoverLedger,)], indirect=True) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_observability.py b/api/tests/unit_tests/services/test_knowledge_fs_observability.py index 7e8a052ce3c..f4c51a9be43 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_observability.py +++ b/api/tests/unit_tests/services/test_knowledge_fs_observability.py @@ -33,12 +33,20 @@ class _FakeInstrument: self.log.append(call) +@dataclass(frozen=True) +class _GaugeRegistration: + name: str + callbacks: tuple[Callable[[object], Iterable[Observation]], ...] + description: str + unit: str + + @dataclass class _FakeMeter: log: list[RecordedCall] = field(default_factory=list) counters: dict[str, _FakeInstrument] = field(default_factory=dict) histograms: dict[str, _FakeInstrument] = field(default_factory=dict) - gauge_calls: list[RecordedCall] = field(default_factory=list) + gauge_calls: list[_GaugeRegistration] = field(default_factory=list) def create_counter(self, name: str, *, description: str, unit: str) -> _FakeInstrument: self.log.append(RecordedCall(args=(name,), kwargs={"description": description, "unit": unit})) @@ -56,9 +64,10 @@ class _FakeMeter: description: str, unit: str, ) -> object: - call = RecordedCall(args=(name,), kwargs={"callbacks": callbacks, "description": description, "unit": unit}) - self.gauge_calls.append(call) - self.log.append(call) + self.gauge_calls.append(_GaugeRegistration(name=name, callbacks=callbacks, description=description, unit=unit)) + self.log.append( + RecordedCall(args=(name,), kwargs={"callbacks": callbacks, "description": description, "unit": unit}) + ) return object() @@ -129,9 +138,9 @@ def test_control_space_current_state_gauge_uses_bounded_states_and_is_best_effor assert len(meter.gauge_calls) == 1 gauge_call = meter.gauge_calls[0] - assert gauge_call.args == ("dify.knowledge_fs.control_spaces",) - assert gauge_call.kwargs["unit"] == "{space}" - callback = gauge_call.kwargs["callbacks"][0] + assert gauge_call.name == "dify.knowledge_fs.control_spaces" + assert gauge_call.unit == "{space}" + callback = gauge_call.callbacks[0] observations = tuple(callback(object())) assert [(observation.value, observation.attributes) for observation in observations] == [ (3, {"aggregation_scope": "global_database_snapshot", "state": "provisioning"}), @@ -139,7 +148,7 @@ def test_control_space_current_state_gauge_uses_bounded_states_and_is_best_effor (1, {"aggregation_scope": "global_database_snapshot", "state": "error"}), ] - assert "aggregate replica series with max, never sum" in gauge_call.kwargs["description"] + assert "aggregate replica series with max, never sum" in gauge_call.description counts_error.append(RuntimeError("database unavailable")) assert tuple(callback(object())) == () diff --git a/api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py b/api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py index 5533caa770b..9b19c68945e 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py +++ b/api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py @@ -9,6 +9,7 @@ from services import knowledge_fs_capability from services.knowledge_fs import lifecycle_remote_http, remote_registry from tests.unit_tests.config_override import apply_config_overrides from tests.unit_tests.knowledge_fs_core_fakes import Marker +from tests.unit_tests.services.knowledge_fs_fakes import UnexpectedLifecycleRemote def test_configured_lifecycle_remote_assembles_capability_v2_http_client( @@ -71,7 +72,7 @@ def test_lifecycle_remote_assembly_fails_closed_without_endpoint_or_capability_i def test_explicit_lifecycle_remote_factory_remains_available_for_isolated_process_tests( monkeypatch: pytest.MonkeyPatch, ) -> None: - remote = Marker("remote") + remote = UnexpectedLifecycleRemote() monkeypatch.setattr(remote_registry, "_remote_factory", None) remote_registry.configure_knowledge_fs_lifecycle_remote_factory(lambda: remote) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_text_store.py b/api/tests/unit_tests/services/test_knowledge_fs_text_store.py index 6653bb1f4ac..495e7bcc6d2 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_text_store.py +++ b/api/tests/unit_tests/services/test_knowledge_fs_text_store.py @@ -21,7 +21,7 @@ from services.knowledge_fs.text_store_native import ( WeaviateTextStore, lexical_sparse_vector, ) -from services.knowledge_fs.vector_store import VectorStoreUnavailableError +from services.knowledge_fs.vector_store import VectorResult, VectorStoreUnavailableError from services.knowledge_fs.vector_store_elasticsearch import ElasticsearchVectorStore from services.knowledge_fs.vector_store_weaviate import WeaviateVectorStore from tests.unit_tests.knowledge_fs_core_fakes import Marker @@ -32,9 +32,10 @@ SCOPE = {"tenant_id": A, "knowledge_space_id": B, "index_version": "fts-v1"} @dataclass(frozen=True) class _StaticTextClient: - result: dict + result: VectorResult - def execute(self, _payload: TextPayload) -> dict: + def execute(self, payload: TextPayload) -> VectorResult: + assert isinstance(payload, TextPayload) return self.result diff --git a/dify-agent/tests/local/dify_agent/agent_stub/server/test_knowledge.py b/dify-agent/tests/local/dify_agent/agent_stub/server/test_knowledge.py index b295270f29a..8895c00403f 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/server/test_knowledge.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/server/test_knowledge.py @@ -4,6 +4,7 @@ import asyncio import base64 import io import json +from dataclasses import dataclass from uuid import uuid4 import httpx @@ -14,6 +15,7 @@ from PIL import Image from pydantic import ValidationError from pydantic_ai.messages import BinaryContent, ModelRequest, ToolReturnPart, UserPromptPart +from dify_agent.adapters.shell.protocols import CompleteShellCommandResult from dify_agent.agent_stub.server.knowledge import AgentStubKnowledgeHandler, _validate_thumbnail, project_result from dify_agent.agent_stub.server.routes.agent_stub import create_agent_stub_http_router from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubPrincipal, AgentStubTokenCodec @@ -34,6 +36,7 @@ from dify_agent.protocol.knowledge_fs import ( KnowledgeFsPreparedRequest, ) from dify_agent.storage.redis_keys import run_cancel_intent_key, run_record_key +from dify_agent.runtime.context import Deps from dify_agent.storage.redis_knowledge_sessions import RedisKnowledgeSessionStore SPACE = "00000000-0000-4000-8000-000000000001" @@ -383,12 +386,31 @@ def test_search_receipt_without_parse_id_reopens_and_pins_provenance(): project_result(command("open", space="docs", receipt_id=RECEIPT), plan, raw, previous) -def knowledge_capability(store, client, *, agent_supports_vision=False, **services): - from types import SimpleNamespace +@dataclass(frozen=True) +class _RunCtx: + """The run-context fields the KnowledgeFS capability reads around one run.""" + deps: Deps + prompt: str + + +def _remote_result(output: str) -> CompleteShellCommandResult: + return CompleteShellCommandResult( + job_id="job", + status="exited", + done=True, + exit_code=0, + output=output, + output_complete=True, + incomplete_reason=None, + offset=len(output), + ) + + +def knowledge_capability(store, client, *, agent_supports_vision=False, **services): from dify_agent.layers.knowledge_fs.layer import Capability, Config from dify_agent.layers.shell.layer import ShellSession - from dify_agent.runtime.context import Deps, Services + from dify_agent.runtime.context import Services deps = Deps( layers={ @@ -411,11 +433,10 @@ def knowledge_capability(store, client, *, agent_supports_vision=False, **servic ) shell = ShellSession("shell", deps) deps.resources.shells["shell"] = shell - return Capability("knowledge_fs"), shell, SimpleNamespace(deps=deps, prompt="inspect evidence") + return Capability("knowledge_fs"), shell, _RunCtx(deps=deps, prompt="inspect evidence") def test_real_layer_delivers_images_as_model_parts_and_cleans_up(monkeypatch): - from types import SimpleNamespace from unittest.mock import AsyncMock from pydantic_ai import Agent from pydantic_ai.models.test import TestModel @@ -425,8 +446,8 @@ def test_real_layer_delivers_images_as_model_parts_and_cleans_up(monkeypatch): redis, store, _ = await session_fixture() remote = AsyncMock( side_effect=[ - SimpleNamespace(exit_code=0, output_complete=True, output="1\n"), - SimpleNamespace(exit_code=0, output_complete=True, output=delivery().result.model_dump_json()), + _remote_result("1\n"), + _remote_result(delivery().result.model_dump_json()), ] ) monkeypatch.setattr(ShellSession, "run_remote_script", remote) @@ -491,7 +512,6 @@ def test_real_layer_delivers_images_as_model_parts_and_cleans_up(monkeypatch): def test_layer_preflight_rejects_incompatible_deployments_and_closes_lease( monkeypatch, protocol, gateway_output, error ): - from types import SimpleNamespace from unittest.mock import AsyncMock from dify_agent.layers.shell.layer import ShellSession @@ -504,8 +524,8 @@ def test_layer_preflight_rejects_incompatible_deployments_and_closes_lease( ).model_dump_json() remote = AsyncMock( side_effect=[ - SimpleNamespace(exit_code=0, output_complete=True, output=protocol), - SimpleNamespace(exit_code=0, output_complete=True, output=output), + _remote_result(protocol), + _remote_result(output), ] ) monkeypatch.setattr(ShellSession, "run_remote_script", remote) diff --git a/dify-agent/tests/local/dify_agent/runtime/test_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_runner.py index 5d3da0a6252..9d901b2bdbf 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_runner.py @@ -504,11 +504,16 @@ async def test_empty_current_user_content_fails_before_model_io(monkeypatch, use @pytest.mark.anyio @pytest.mark.parametrize("failed", [False, True]) async def test_quality_report_follows_final_run_event(monkeypatch, failed): + from dataclasses import dataclass from unittest.mock import AsyncMock from dify_agent.runtime.runner import RunSuccessOutcome - module = SimpleNamespace(publish_investigation=AsyncMock()) + @dataclass(frozen=True) + class QualityReportModule: + publish_investigation: AsyncMock + + module = QualityReportModule(publish_investigation=AsyncMock()) async with httpx.AsyncClient() as client: services = Services(client, client) runner = AgentRunRunner(sink=InMemoryRunEventSink(), request=request(), run_id="quality-run", services=services)