diff --git a/api/controllers/console/app/mcp_server.py b/api/controllers/console/app/mcp_server.py index 59c851d1c8c..d03e27b7fae 100644 --- a/api/controllers/console/app/mcp_server.py +++ b/api/controllers/console/app/mcp_server.py @@ -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() diff --git a/api/controllers/console/auth/data_source_bearer_auth.py b/api/controllers/console/auth/data_source_bearer_auth.py index fac725e8534..3f42b7a3ad7 100644 --- a/api/controllers/console/auth/data_source_bearer_auth.py +++ b/api/controllers/console/auth/data_source_bearer_auth.py @@ -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()) diff --git a/api/controllers/console/tag/tags.py b/api/controllers/console/tag/tags.py index 86c1ad9c54c..00596545d1c 100644 --- a/api/controllers/console/tag/tags.py +++ b/api/controllers/console/tag/tags.py @@ -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()) diff --git a/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py b/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py index a6568d5932a..0978a79b27c 100644 --- a/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py +++ b/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py @@ -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" ), diff --git a/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py b/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py index b231826aeac..549a78a6915 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py +++ b/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py @@ -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"} diff --git a/api/tests/unit_tests/controllers/console/tag/test_tags.py b/api/tests/unit_tests/controllers/console/tag/test_tags.py index 7d267b50153..55adfab959f 100644 --- a/api/tests/unit_tests/controllers/console/tag/test_tags.py +++ b/api/tests/unit_tests/controllers/console/tag/test_tags.py @@ -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()