Kimirock/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test: use real tool model and encryption fixtures (#43365)

Browse Source
This commit is contained in:
Asuka Minato
2026-10-04 11:38:24 +00:00
committed by GitHub
parent 16bbfc50bd
commit f0266537cb
2 changed files with 87 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
api/tests/unit_tests/core/tools/utils/test_encryption.py
View File

@@ -1,31 +1,17 @@
import copy
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock, patch
from unittest.mock import patch
import pytest
from core.entities.provider_entities import BasicProviderConfig, ProviderConfigType
from core.entities.provider_entities import BasicProviderConfig, ProviderConfig, ProviderConfigType
from core.helper.provider_cache import NoOpProviderCredentialCache, SingletonProviderCredentialsCache
from core.helper.provider_encryption import ProviderConfigEncrypter
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_entities import ToolProviderEntityWithPlugin, ToolProviderIdentity
from core.tools.plugin_tool.provider import PluginToolProviderController
from core.tools.utils.encryption import create_tool_provider_encrypter
# ---------------------------
# A no-op cache
# ---------------------------
class NoopCache:
"""Simple cache stub: always returns None, does nothing for set/delete."""
def get(self) -> Any | None:
return None
def set(self, config: Any) -> None:
pass
def delete(self) -> None:
pass
@pytest.fixture
def secret_field() -> BasicProviderConfig:
"""A SECRET_INPUT field named 'password'."""
@@ -50,12 +36,12 @@ def encrypter_obj(secret_field, normal_field):
Build ProviderConfigEncrypter with:
- tenant_id = tenant123
- one secret field (password) and one normal field (username)
- NoopCache as cache
- NoOpProviderCredentialCache as cache
"""
return ProviderConfigEncrypter(
tenant_id="tenant123",
config=[secret_field, normal_field],
provider_config_cache=NoopCache(),
provider_config_cache=NoOpProviderCredentialCache(),
)
@@ -186,31 +172,30 @@ def test_decrypt_swallow_exception_and_keep_original(encrypter_obj):
def test_create_tool_provider_encrypter_builds_cache_and_encrypter():
basic_config = BasicProviderConfig(name="key", type=ProviderConfigType.TEXT_INPUT)
credential_schema_item = SimpleNamespace(to_basic_provider_config=lambda: basic_config)
controller = SimpleNamespace(
provider_type=SimpleNamespace(value="builtin"),
entity=SimpleNamespace(identity=SimpleNamespace(name="provider-a")),
get_credentials_schema=lambda: [credential_schema_item],
)
cache_instance = Mock()
encrypter_instance = Mock()
with patch(
"core.tools.utils.encryption.SingletonProviderCredentialsCache", return_value=cache_instance
) as cache_cls:
with patch("core.tools.utils.encryption.ProviderConfigEncrypter", return_value=encrypter_instance) as enc_cls:
encrypter, cache = create_tool_provider_encrypter("tenant-1", controller)
assert encrypter is encrypter_instance
assert cache is cache_instance
cache_cls.assert_called_once_with(
controller = PluginToolProviderController(
entity=ToolProviderEntityWithPlugin(
identity=ToolProviderIdentity(
author="author",
name="provider-a",
description=I18nObject(en_US="Description"),
icon="icon.svg",
label=I18nObject(en_US="Provider"),
),
credentials_schema=[ProviderConfig(name="key", type=ProviderConfigType.TEXT_INPUT)],
plugin_id="plugin-id",
tools=[],
),
plugin_id="plugin-id",
plugin_unique_identifier="plugin-uid",
tenant_id="tenant-1",
provider_type="builtin",
provider_identity="provider-a",
)
enc_cls.assert_called_once_with(
tenant_id="tenant-1",
config=[basic_config],
provider_config_cache=cache_instance,
)
encrypter, cache = create_tool_provider_encrypter("tenant-1", controller)
assert isinstance(encrypter, ProviderConfigEncrypter)
assert isinstance(cache, SingletonProviderCredentialsCache)
assert encrypter.tenant_id == "tenant-1"
assert encrypter.config == [basic_config]
assert encrypter.provider_config_cache is cache
assert cache.cache_key == "plugin_credentials:tenant_id:tenant-1:id:plugin.provider-a"
assert encrypter.encrypt({"key": "value"}) == {"key": "value"}

96
api/tests/unit_tests/core/tools/utils/test_model_invocation_utils.py
View File

@@ -3,21 +3,21 @@
Covers success and error branches for ModelInvocationUtils, including
InvokeModelError and invoke error mappings for InvokeAuthorizationError,
InvokeBadRequestError, InvokeConnectionError, InvokeRateLimitError, and
InvokeServerUnavailableError. Model-provider collaborators remain mocked, while
invocation logging uses real SQLite-backed SQLAlchemy sessions.
InvokeServerUnavailableError. Real model instances supply validated schemas and
results, while invocation logging uses real SQLite-backed SQLAlchemy sessions.
"""
from __future__ import annotations
from decimal import Decimal
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock, patch
from unittest.mock import patch
import pytest
from sqlalchemy import event, select
from sqlalchemy.orm import Session, sessionmaker
from core.model_manager import ModelManager
from core.tools.entities.tool_entities import ToolProviderType
from core.tools.utils.model_invocation_utils import InvokeModelError, ModelInvocationUtils
from graphon.model_runtime.entities.model_entities import ModelPropertyKey
@@ -29,35 +29,20 @@ from graphon.model_runtime.errors.invoke import (
InvokeServerUnavailableError,
)
from models.tools import ToolModelInvoke
from tests.unit_tests.core.model_fixtures import make_model_config, make_model_instance
TENANT_ID = "11111111-1111-1111-1111-111111111111"
USER_ID = "22222222-2222-2222-2222-222222222222"
CALLER_ID = "33333333-3333-3333-3333-333333333333"
def _mock_model_instance(*, schema: dict[str, Any] | None = None) -> SimpleNamespace:
model_type_instance = Mock()
model_type_instance.get_model_schema.return_value = (
SimpleNamespace(model_properties=schema or {}) if schema is not None else None
)
return SimpleNamespace(
provider="provider",
model="model-a",
model_name="model-a",
credentials={"api_key": "x"},
model_type_instance=model_type_instance,
get_llm_num_tokens=lambda prompt_messages: 5,
invoke_llm=Mock(),
)
@pytest.mark.parametrize(
("model_instance", "expected", "error_match"),
("model_exists", "properties", "expected", "error_match"),
[
(None, None, "Model not found"),
(_mock_model_instance(schema=None), None, "No model schema found"),
(_mock_model_instance(schema={}), 2048, None),
(_mock_model_instance(schema={ModelPropertyKey.CONTEXT_SIZE: 8192}), 8192, None),
(False, None, None, "Model not found"),
(True, None, None, "No model schema found"),
(True, {}, 2048, None),
(True, {ModelPropertyKey.CONTEXT_SIZE: 8192}, 8192, None),
],
ids=[
"missing-model",
@@ -66,11 +51,22 @@ def _mock_model_instance(*, schema: dict[str, Any] | None = None) -> SimpleNames
"schema-context-size",
],
)
def test_get_max_llm_context_tokens_branches(model_instance, expected, error_match):
manager = Mock()
manager.get_default_model_instance.return_value = model_instance
def test_get_max_llm_context_tokens_branches(model_exists, properties, expected, error_match):
manager = ModelManager.for_tenant("tenant", user_id="user-1")
model_instance = make_model_instance(provider="provider", model="model-a")
schema = make_model_config(provider="provider", model="model-a", mode="chat").model_schema
if properties is not None:
schema.model_properties = properties
with patch("core.tools.utils.model_invocation_utils.ModelManager.for_tenant", return_value=manager) as mock_factory:
with (
patch("core.tools.utils.model_invocation_utils.ModelManager.for_tenant", return_value=manager) as mock_factory,
patch.object(manager, "get_default_model_instance", return_value=model_instance if model_exists else None),
patch.object(
model_instance.model_type_instance,
"get_model_schema",
return_value=schema if properties is not None else None,
),
):
if error_match:
with pytest.raises(InvokeModelError, match=error_match):
ModelInvocationUtils.get_max_llm_context_tokens("tenant", user_id="user-1")
@@ -81,29 +77,37 @@ def test_get_max_llm_context_tokens_branches(model_instance, expected, error_mat
def test_calculate_tokens_handles_missing_model():
manager = Mock()
manager.get_default_model_instance.return_value = None
with patch("core.tools.utils.model_invocation_utils.ModelManager.for_tenant", return_value=manager) as mock_factory:
manager = ModelManager.for_tenant("tenant")
with (
patch("core.tools.utils.model_invocation_utils.ModelManager.for_tenant", return_value=manager) as mock_factory,
patch.object(manager, "get_default_model_instance", return_value=None),
):
with pytest.raises(InvokeModelError, match="Model not found"):
ModelInvocationUtils.calculate_tokens("tenant", [])
mock_factory.assert_called_once_with(tenant_id="tenant", user_id=None)
def test_invoke_success_and_error_mappings(sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]):
model_instance = _mock_model_instance(schema={ModelPropertyKey.CONTEXT_SIZE: 2048})
model_instance.invoke_llm.return_value = SimpleNamespace(
message=SimpleNamespace(content="ok"),
usage=SimpleNamespace(
model_instance = make_model_instance(provider="provider", model="model-a")
result = LLMResult(
model="model-a",
message=AssistantPromptMessage(content="ok"),
usage=LLMUsage(
prompt_tokens=5,
prompt_unit_price=Decimal(0),
prompt_price_unit=Decimal(1),
prompt_price=Decimal(0),
completion_tokens=7,
completion_unit_price=Decimal("0.1"),
completion_price_unit=Decimal(1),
completion_price=Decimal("0.7"),
total_tokens=12,
latency=0.3,
total_price=Decimal("0.7"),
currency="USD",
),
)
manager = Mock()
manager.get_default_model_instance.return_value = model_instance
manager = ModelManager.for_tenant(TENANT_ID, user_id=CALLER_ID)
database = SimpleNamespace(session=sqlite_session)
attached_invocations: list[ToolModelInvoke] = []
@@ -122,6 +126,9 @@ def test_invoke_success_and_error_mappings(sqlite_session: Session, sqlite_sessi
with (
patch("core.tools.utils.model_invocation_utils.ModelManager.for_tenant", return_value=manager) as mock_factory,
patch.object(manager, "get_default_model_instance", return_value=model_instance),
patch.object(model_instance, "get_llm_num_tokens", return_value=5),
patch.object(model_instance, "invoke_llm", return_value=result),
patch("core.tools.utils.model_invocation_utils.db", database),
):
response = ModelInvocationUtils.invoke(
@@ -170,15 +177,16 @@ def test_invoke_success_and_error_mappings(sqlite_session: Session, sqlite_sessi
],
)
def test_invoke_error_mappings(exc, expected, sqlite_session: Session):
model_instance = _mock_model_instance(schema={ModelPropertyKey.CONTEXT_SIZE: 2048})
model_instance.invoke_llm.side_effect = exc
manager = Mock()
manager.get_default_model_instance.return_value = model_instance
model_instance = make_model_instance(provider="provider", model="model-a")
manager = ModelManager.for_tenant(TENANT_ID, user_id=USER_ID)
database = SimpleNamespace(session=sqlite_session)
with (
patch("core.tools.utils.model_invocation_utils.ModelManager.for_tenant", return_value=manager) as mock_factory,
patch.object(manager, "get_default_model_instance", return_value=model_instance),
patch.object(model_instance, "get_llm_num_tokens", return_value=5),
patch.object(model_instance, "invoke_llm", side_effect=exc),
patch("core.tools.utils.model_invocation_utils.db", database),
):
with pytest.raises(InvokeModelError, match=expected):
@@ -196,3 +204,7 @@ def test_invoke_error_mappings(exc, expected, sqlite_session: Session):
assert persisted.prompt_tokens == 5
assert persisted.answer_tokens == 0
mock_factory.assert_called_once_with(tenant_id=TENANT_ID, user_id=USER_ID)
from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage