opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-01 15:09:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: dep inject for model in tags, mcp server, and data source auth (#40159)

Browse Source
This commit is contained in:
Likalikali
2026-08-07 20:38:29 +08:00
committed by GitHub
parent 06e229e8b2
commit 055462fe6e
6 changed files with 59 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/controllers/console/app/mcp_server.py
+12 -12
View File
@@ -15,6 +15,7 @@ from controllers.console.wraps import (
RBACResourceScope,
account_initialization_required,
edit_permission_required,
model_validate,
rbac_permission_required,
setup_required,
with_current_tenant_id,
@@ -109,17 +110,16 @@ class AppMCPServerController(Resource):
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@with_current_tenant_id
@get_app_model
def post(self, current_tenant_id: str, app_model: App):
payload = MCPServerCreatePayload.model_validate(console_ns.payload or {})
description = payload.description
@model_validate(MCPServerCreatePayload)
def post(self, req_data: MCPServerCreatePayload, current_tenant_id: str, app_model: App):
description = req_data.description
if not description:
description = app_model.description or ""
server = AppMCPServer(
name=app_model.name,
description=description,
parameters=json.dumps(payload.parameters, ensure_ascii=False),
parameters=json.dumps(req_data.parameters, ensure_ascii=False),
status=AppMCPServerStatus.ACTIVE,
app_id=app_model.id,
tenant_id=current_tenant_id,
@@ -144,10 +144,10 @@ class AppMCPServerController(Resource):
@edit_permission_required
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
@get_app_model
def put(self, app_model: App):
payload = MCPServerUpdatePayload.model_validate(console_ns.payload or {})
@model_validate(MCPServerUpdatePayload)
def put(self, req_data: MCPServerUpdatePayload, app_model: App):
app_ref = AppRefService.create_app_ref(app_model)
server_ref = AppRefService.create_mcp_server_ref(app_ref, payload.id)
server_ref = AppRefService.create_mcp_server_ref(app_ref, req_data.id)
server = db.session.scalar(
select(AppMCPServer)
.where(
@@ -160,7 +160,7 @@ class AppMCPServerController(Resource):
if not server:
raise NotFound()
description = payload.description
description = req_data.description
if description is None or not description:
server.description = app_model.description or ""
else:
@@ -168,10 +168,10 @@ class AppMCPServerController(Resource):
server.name = app_model.name
server.parameters = json.dumps(payload.parameters, ensure_ascii=False)
if payload.status:
server.parameters = json.dumps(req_data.parameters, ensure_ascii=False)
if req_data.status:
try:
server.status = AppMCPServerStatus(payload.status)
server.status = AppMCPServerStatus(req_data.status)
except ValueError:
raise ValueError("Invalid status")
db.session.commit()
api/controllers/console/auth/data_source_bearer_auth.py
+4 -3
View File
@@ -17,6 +17,7 @@ from ..wraps import (
RBACResourceScope,
account_initialization_required,
is_admin_or_owner_required,
model_validate,
rbac_permission_required,
setup_required,
with_current_tenant_id,
@@ -87,10 +88,10 @@ class ApiKeyAuthDataSourceBinding(Resource):
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_CREATE, resource_required=False)
@console_ns.expect(console_ns.models[ApiKeyAuthBindingPayload.__name__])
@with_current_tenant_id
def post(self, current_tenant_id: str):
@model_validate(ApiKeyAuthBindingPayload)
def post(self, req_data: ApiKeyAuthBindingPayload, current_tenant_id: str):
# The role of the current user in the table must be admin or owner
payload = ApiKeyAuthBindingPayload.model_validate(console_ns.payload)
data = payload.model_dump()
data = req_data.model_dump()
ApiKeyAuthService.validate_api_key_auth_args(data)
try:
ApiKeyAuthService.create_provider_auth(current_tenant_id, data, session=db.session())
api/controllers/console/tag/tags.py
+8 -7
View File
@@ -17,6 +17,7 @@ from controllers.console.wraps import (
RBACResourceScope,
account_initialization_required,
edit_permission_required,
model_validate,
setup_required,
with_current_tenant_id,
with_current_user,
@@ -147,14 +148,14 @@ class TagListApi(Resource):
@login_required
@account_initialization_required
@with_current_user
def post(self, current_user: Account):
@model_validate(TagBasePayload)
def post(self, req_data: TagBasePayload, current_user: Account):
# Allow users with edit permission, or dataset editors (including dataset operators).
if not (current_user.has_edit_permission or current_user.is_dataset_editor):
raise Forbidden()
payload = TagBasePayload.model_validate(console_ns.payload or {})
_enforce_snippet_tag_rbac_if_needed(payload.type)
tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=payload.type), db.session())
_enforce_snippet_tag_rbac_if_needed(req_data.type)
tag = TagService.save_tags(SaveTagPayload(name=req_data.name, type=req_data.type), db.session())
return dump_response(TagResponse, {"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": 0}), 200
@@ -167,15 +168,15 @@ class TagUpdateDeleteApi(Resource):
@login_required
@account_initialization_required
@with_current_user
def patch(self, current_user: Account, tag_id: UUID):
@model_validate(TagUpdateRequestPayload)
def patch(self, req_data: TagUpdateRequestPayload, current_user: Account, tag_id: UUID):
tag_id_str = str(tag_id)
# The role of the current user in the ta table must be admin, owner, or editor
if not (current_user.has_edit_permission or current_user.is_dataset_editor):
raise Forbidden()
payload = TagUpdateRequestPayload.model_validate(console_ns.payload or {})
_enforce_snippet_tag_rbac_by_tag_id(tag_id_str)
tag = TagService.update_tags(UpdateTagPayload(name=payload.name), tag_id_str, db.session())
tag = TagService.update_tags(UpdateTagPayload(name=req_data.name), tag_id_str, db.session())
binding_count = TagService.get_tag_binding_count(tag_id_str, db.session())
api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py
+12 -5
View File
@@ -1,7 +1,7 @@
import datetime
from inspect import unwrap
from types import SimpleNamespace
from unittest.mock import PropertyMock, patch
from unittest.mock import patch
import pytest
from flask import Flask
@@ -14,6 +14,8 @@ from controllers.console.app.mcp_server import (
AppMCPServerController,
AppMCPServerRefreshController,
AppMCPServerResponse,
MCPServerCreatePayload,
MCPServerUpdatePayload,
)
from controllers.console.wraps import RBACPermission, RBACResourceScope
from models.enums import AppMCPServerStatus
@@ -142,17 +144,20 @@ class TestAppMCPServerController:
api = AppMCPServerController()
method = unwrap(api.post)
payload = {"parameters": {"timeout": 30}}
req_data = MCPServerCreatePayload.model_validate(payload)
app = Flask(__name__)
app.config["TESTING"] = True
with (
app.test_request_context("/", json=payload),
patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
patch("controllers.console.app.mcp_server.db.session", sqlite_session),
patch("controllers.console.app.mcp_server.AppMCPServer.generate_server_code", return_value="server-code"),
):
response, status_code = method(
api, "tenant-1", app_model=SimpleNamespace(id="app-1", name="Demo App", description="App description")
api,
req_data,
"tenant-1",
app_model=SimpleNamespace(id="app-1", name="Demo App", description="App description"),
)
server = sqlite_session.scalar(select(AppMCPServer))
@@ -165,6 +170,7 @@ class TestAppMCPServerController:
api = AppMCPServerController()
method = unwrap(api.put)
payload = {"id": "server-1", "description": "Updated", "parameters": {"timeout": 30}, "status": "active"}
req_data = MCPServerUpdatePayload.model_validate(payload)
app = Flask(__name__)
app.config["TESTING"] = True
server = _server(name="Old", description="Old")
@@ -174,11 +180,11 @@ class TestAppMCPServerController:
with (
app.test_request_context("/", json=payload),
patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
patch("controllers.console.app.mcp_server.db.session", sqlite_session),
):
response = method(
api,
req_data,
app_model=SimpleNamespace(
id="app-1", tenant_id="tenant-1", name="Demo App", description="App description"
),
@@ -206,6 +212,7 @@ class TestAppMCPServerController:
api = AppMCPServerController()
method = unwrap(api.put)
payload = {"id": "server-1", "description": "Updated", "parameters": {"timeout": 30}, "status": "active"}
req_data = MCPServerUpdatePayload.model_validate(payload)
app = Flask(__name__)
app.config["TESTING"] = True
foreign_server = _server(
@@ -220,12 +227,12 @@ class TestAppMCPServerController:
with (
app.test_request_context("/", json=payload),
patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
patch("controllers.console.app.mcp_server.db.session", sqlite_session),
pytest.raises(NotFound),
):
method(
api,
req_data,
app_model=SimpleNamespace(
id="app-1", tenant_id="tenant-1", name="Demo App", description="App description"
),
api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py
+4 -13
View File
@@ -3,25 +3,16 @@ from __future__ import annotations
from datetime import UTC, datetime
from inspect import unwrap
from types import SimpleNamespace
from unittest.mock import ANY, PropertyMock, patch
from unittest.mock import ANY, patch
from controllers.console import console_ns
from controllers.console.auth.data_source_bearer_auth import (
ApiKeyAuthBindingPayload,
ApiKeyAuthDataSource,
ApiKeyAuthDataSourceBinding,
ApiKeyAuthDataSourceBindingDelete,
)
def _payload_patch(payload: dict):
return patch.object(
type(console_ns),
"payload",
new_callable=PropertyMock,
return_value=payload,
)
def test_list_data_source_auth_uses_injected_tenant_id() -> None:
api = ApiKeyAuthDataSource()
method = unwrap(api.get)
@@ -56,14 +47,14 @@ def test_create_data_source_auth_binding_uses_injected_tenant_id() -> None:
"provider": "custom",
"credentials": {"auth_type": "api_key", "config": {"api_key": "secret"}},
}
req_data = ApiKeyAuthBindingPayload.model_validate(payload)
with (
_payload_patch(payload),
patch("controllers.console.auth.data_source_bearer_auth.db"),
patch("controllers.console.auth.data_source_bearer_auth.ApiKeyAuthService.validate_api_key_auth_args"),
patch("controllers.console.auth.data_source_bearer_auth.ApiKeyAuthService.create_provider_auth") as create_auth,
):
result, status = method(api, "tenant-1")
result, status = method(api, req_data, "tenant-1")
create_auth.assert_called_once_with("tenant-1", payload, session=ANY)
assert result == {"result": "success"}
api/tests/unit_tests/controllers/console/tag/test_tags.py
+19 -27
View File
@@ -11,10 +11,12 @@ from werkzeug.exceptions import Forbidden
import controllers.console.tag.tags as module
from controllers.console import console_ns
from controllers.console.tag.tags import (
TagBasePayload,
TagBindingCollectionApi,
TagBindingRemoveApi,
TagListApi,
TagUpdateDeleteApi,
TagUpdateRequestPayload,
)
from models import Account
from models.account import AccountStatus, TenantAccountRole
@@ -161,35 +163,35 @@ class TestTagListApi:
assert status == 200
assert result == [{"id": "1", "name": "snippet-tag", "type": "snippet", "binding_count": "1"}]
def test_post_success(self, app: Flask, admin_user, tag, payload_patch):
def test_post_success(self, app: Flask, admin_user, tag):
api = TagListApi()
method = unwrap(api.post)
payload = {"name": "test-tag", "type": "knowledge"}
req_data = TagBasePayload.model_validate(payload)
with app.test_request_context("/", json=payload):
with (
payload_patch(payload),
patch(
"controllers.console.tag.tags.TagService.save_tags",
return_value=tag,
),
):
result, status = method(api, admin_user)
result, status = method(api, req_data, admin_user)
assert status == 200
assert result["name"] == "test-tag"
assert result["binding_count"] == "0"
def test_post_snippet_tag_checks_snippet_rbac_when_enabled(self, app: Flask, admin_user, tag, payload_patch):
def test_post_snippet_tag_checks_snippet_rbac_when_enabled(self, app: Flask, admin_user, tag):
api = TagListApi()
method = unwrap(api.post)
payload = {"name": "snippet-tag", "type": "snippet"}
req_data = TagBasePayload.model_validate(payload)
with app.test_request_context("/", json=payload):
with (
payload_patch(payload),
patch("controllers.console.tag.tags.dify_config.RBAC_ENABLED", True),
patch(
"controllers.console.tag.tags.current_account_with_tenant",
@@ -201,7 +203,7 @@ class TestTagListApi:
return_value=tag,
),
):
method(api, admin_user)
method(api, req_data, admin_user)
enforce_mock.assert_called_once_with(
tenant_id="tenant-1",
@@ -211,30 +213,25 @@ class TestTagListApi:
resource_required=False,
)
def test_post_forbidden(self, app: Flask, readonly_user, payload_patch):
def test_post_forbidden(self, app: Flask, readonly_user):
api = TagListApi()
method = unwrap(api.post)
payload = {"name": "x"}
with app.test_request_context("/", json=payload):
with (
payload_patch(payload),
):
with pytest.raises(Forbidden):
method(api, readonly_user)
with app.test_request_context("/"):
with pytest.raises(Forbidden):
method(api, None, readonly_user)
class TestTagUpdateDeleteApi:
def test_patch_success(self, app: Flask, admin_user, tag, payload_patch, sqlite_engine: Engine):
def test_patch_success(self, app: Flask, admin_user, tag, sqlite_engine: Engine):
api = TagUpdateDeleteApi()
method = unwrap(api.patch)
payload = {"name": "updated"}
req_data = TagUpdateRequestPayload.model_validate(payload)
with app.test_request_context("/", json=payload):
with (
payload_patch(payload),
patch(
"controllers.console.tag.tags.TagService.update_tags",
return_value=tag,
@@ -244,7 +241,7 @@ class TestTagUpdateDeleteApi:
return_value=3,
),
):
result, status = method(api, admin_user, "tag-1")
result, status = method(api, req_data, admin_user, "tag-1")
assert status == 200
update_payload, tag_id, session = update_tags_mock.call_args.args
@@ -253,18 +250,13 @@ class TestTagUpdateDeleteApi:
_assert_sqlite_session(session, sqlite_engine)
assert result["binding_count"] == "3"
def test_patch_forbidden(self, app: Flask, readonly_user, payload_patch):
def test_patch_forbidden(self, app: Flask, readonly_user):
api = TagUpdateDeleteApi()
method = unwrap(api.patch)
payload = {"name": "x"}
with app.test_request_context("/", json=payload):
with (
payload_patch(payload),
):
with pytest.raises(Forbidden):
method(api, readonly_user, "tag-1")
with app.test_request_context("/"):
with pytest.raises(Forbidden):
method(api, None, readonly_user, "tag-1")
def test_delete_success(self, app: Flask, admin_user, sqlite_engine: Engine):
api = TagUpdateDeleteApi()