From a1577fc48c8fec17fcb6873a43b8c2155e59afdf Mon Sep 17 00:00:00 2001 From: Likalikali Date: Sun, 9 Aug 2026 14:14:07 +0800 Subject: [PATCH] refactor: replace manual model_validate with @model_validate in auth/billing/tag controllers (#40240) --- api/controllers/console/auth/activate.py | 25 +++++----- .../console/auth/email_register.py | 44 ++++++++--------- .../console/auth/forgot_password.py | 38 +++++++-------- api/controllers/console/auth/login.py | 47 ++++++++++--------- api/controllers/console/billing/billing.py | 14 +++--- api/controllers/console/billing/compliance.py | 7 +-- api/controllers/console/remote_files.py | 2 +- api/controllers/console/tag/tags.py | 24 +++++----- .../console/billing/test_billing.py | 8 ++-- .../controllers/console/tag/test_tags.py | 29 ++++++++---- 10 files changed, 125 insertions(+), 113 deletions(-) diff --git a/api/controllers/console/auth/activate.py b/api/controllers/console/auth/activate.py index 3e9160f2bb0..d29e43f7f09 100644 --- a/api/controllers/console/auth/activate.py +++ b/api/controllers/console/auth/activate.py @@ -9,6 +9,7 @@ from controllers.common.schema import query_params_from_model, register_schema_m from controllers.console import console_ns from controllers.console.auth.error import InvitationAccountMismatchError from controllers.console.error import AccountInFreezeError, AlreadyActivateError +from controllers.console.wraps import model_validate from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from libs.helper import EmailStr, timezone @@ -86,14 +87,14 @@ class ActivateCheckApi(Resource): "Success", console_ns.models[ActivationCheckResponse.__name__], ) - def get(self): - args = ActivateCheckQuery.model_validate(request.args.to_dict(flat=True)) + @model_validate(ActivateCheckQuery) + def get(self, req_data: ActivateCheckQuery): - workspaceId = args.workspace_id - token = args.token + workspaceId = req_data.workspace_id + token = req_data.token invitation = RegisterService.get_invitation_with_case_fallback( - workspaceId, args.email, token, session=db.session() + workspaceId, req_data.email, token, session=db.session() ) if invitation: data = invitation.get("data", {}) @@ -138,18 +139,18 @@ class ActivateApi(Resource): console_ns.models[ActivationResponse.__name__], ) @console_ns.response(400, "Already activated or invalid token") - def post(self): + @model_validate(ActivatePayload) + def post(self, req_data: ActivatePayload): """Accept an invitation without letting an existing session act for another account. Token-only activation remains available for legacy clients. When the request already carries a console session, that session must belong to the account encoded in the invitation before the token is consumed or tenant membership is changed. """ - args = ActivatePayload.model_validate(console_ns.payload) - normalized_request_email = args.email.lower() if args.email else None + normalized_request_email = req_data.email.lower() if req_data.email else None invitation = RegisterService.get_invitation_with_case_fallback( - args.workspace_id, args.email, args.token, session=db.session() + req_data.workspace_id, req_data.email, req_data.token, session=db.session() ) if invitation is None: raise AlreadyActivateError() @@ -185,11 +186,11 @@ class ActivateApi(Resource): setup_fields: tuple[str, str, str] | None = None if requires_setup: - if not args.name or not args.interface_language or not args.timezone: + if not req_data.name or not req_data.interface_language or not req_data.timezone: raise AlreadyActivateError() - setup_fields = (args.name, args.interface_language, args.timezone) + setup_fields = (req_data.name, req_data.interface_language, req_data.timezone) - RegisterService.revoke_token(args.workspace_id, normalized_request_email, args.token) + RegisterService.revoke_token(req_data.workspace_id, normalized_request_email, req_data.token) if membership_id is None: TenantService.create_tenant_member(tenant, account, db.session(), role=role) diff --git a/api/controllers/console/auth/email_register.py b/api/controllers/console/auth/email_register.py index a9b73ca4679..6ee65e87344 100644 --- a/api/controllers/console/auth/email_register.py +++ b/api/controllers/console/auth/email_register.py @@ -26,7 +26,7 @@ from services.billing_service import BillingService from services.errors.account import AccountRegisterError, SeatsLimitExceededError from ..error import AccountInFreezeError, EmailSendIpLimitError, SeatsLimitExceeded -from ..wraps import email_password_login_enabled, email_register_enabled, setup_required +from ..wraps import email_password_login_enabled, email_register_enabled, model_validate, setup_required class EmailRegisterSendPayload(BaseModel): @@ -87,21 +87,21 @@ class EmailRegisterSendEmailApi(Resource): @email_register_enabled @console_ns.expect(console_ns.models[EmailRegisterSendPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultDataResponse.__name__]) - def post(self): - args = EmailRegisterSendPayload.model_validate(console_ns.payload) - normalized_email = args.email.lower() + @model_validate(EmailRegisterSendPayload) + def post(self, req_data: EmailRegisterSendPayload): + normalized_email = req_data.email.lower() ip_address = extract_remote_ip(request) if AccountService.is_email_send_ip_limit(ip_address): raise EmailSendIpLimitError() language = "en-US" - if args.language is not None and args.language in languages: - language = args.language + if req_data.language is not None and req_data.language in languages: + language = req_data.language if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(normalized_email): raise AccountInFreezeError() - account = AccountService.get_account_by_email_with_case_fallback(args.email, session=db.session()) + account = AccountService.get_account_by_email_with_case_fallback(req_data.email, session=db.session()) token = AccountService.send_email_register_email(email=normalized_email, account=account, language=language) return {"result": "success", "data": token} @@ -113,16 +113,16 @@ class EmailRegisterCheckApi(Resource): @email_register_enabled @console_ns.expect(console_ns.models[EmailRegisterValidityPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[VerificationTokenResponse.__name__]) - def post(self): - args = EmailRegisterValidityPayload.model_validate(console_ns.payload) + @model_validate(EmailRegisterValidityPayload) + def post(self, req_data: EmailRegisterValidityPayload): - user_email = args.email.lower() + user_email = req_data.email.lower() is_email_register_error_rate_limit = AccountService.is_email_register_error_rate_limit(user_email) if is_email_register_error_rate_limit: raise EmailRegisterLimitError() - token_data = AccountService.get_email_register_data(args.token) + token_data = AccountService.get_email_register_data(req_data.token) if token_data is None: raise InvalidTokenError() @@ -132,16 +132,16 @@ class EmailRegisterCheckApi(Resource): if user_email != normalized_token_email: raise InvalidEmailError() - if args.code != token_data.get("code"): + if req_data.code != token_data.get("code"): AccountService.add_email_register_error_rate_limit(user_email) raise EmailCodeError() # Verified, revoke the first token - AccountService.revoke_email_register_token(args.token) + AccountService.revoke_email_register_token(req_data.token) # Refresh token data by generating a new token _, new_token = AccountService.generate_email_register_token( - user_email, code=args.code, additional_data={"phase": "register"} + user_email, code=req_data.code, additional_data={"phase": "register"} ) AccountService.reset_email_register_error_rate_limit(user_email) @@ -155,15 +155,15 @@ class EmailRegisterResetApi(Resource): @email_register_enabled @console_ns.expect(console_ns.models[EmailRegisterResetPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[EmailRegisterResetResponse.__name__]) - def post(self): - args = EmailRegisterResetPayload.model_validate(console_ns.payload) + @model_validate(EmailRegisterResetPayload) + def post(self, req_data: EmailRegisterResetPayload): # Validate passwords match - if args.new_password != args.password_confirm: + if req_data.new_password != req_data.password_confirm: raise PasswordMismatchError() # Validate token and get register data - register_data = AccountService.get_email_register_data(args.token) + register_data = AccountService.get_email_register_data(req_data.token) if not register_data: raise InvalidTokenError() # Must use token in reset phase @@ -171,7 +171,7 @@ class EmailRegisterResetApi(Resource): raise InvalidTokenError() # Revoke token to prevent reuse - AccountService.revoke_email_register_token(args.token) + AccountService.revoke_email_register_token(req_data.token) email = register_data.get("email", "") normalized_email = email.lower() @@ -183,9 +183,9 @@ class EmailRegisterResetApi(Resource): account = self._create_new_account( email=normalized_email, - password=args.password_confirm, - timezone=args.timezone, - language=args.language, + password=req_data.password_confirm, + timezone=req_data.timezone, + language=req_data.language, ) token_pair = AccountService.login(account=account, session=db.session(), ip_address=extract_remote_ip(request)) AccountService.reset_login_error_rate_limit(normalized_email) diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index 8a46a2559cf..9a8784d543a 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/api/controllers/console/auth/forgot_password.py @@ -15,7 +15,7 @@ from controllers.console.auth.error import ( PasswordMismatchError, ) from controllers.console.error import AccountNotFound, EmailSendIpLimitError -from controllers.console.wraps import email_password_login_enabled, setup_required +from controllers.console.wraps import email_password_login_enabled, model_validate, setup_required from extensions.ext_database import db from libs.helper import EmailStr, extract_remote_ip from libs.password import hash_password @@ -68,20 +68,20 @@ class ForgotPasswordSendEmailApi(Resource): @console_ns.response(400, "Invalid email or rate limit exceeded") @setup_required @email_password_login_enabled - def post(self): - args = ForgotPasswordSendPayload.model_validate(console_ns.payload) - normalized_email = args.email.lower() + @model_validate(ForgotPasswordSendPayload) + def post(self, req_data: ForgotPasswordSendPayload): + normalized_email = req_data.email.lower() ip_address = extract_remote_ip(request) if AccountService.is_email_send_ip_limit(ip_address): raise EmailSendIpLimitError() - if args.language is not None and args.language == "zh-Hans": + if req_data.language is not None and req_data.language == "zh-Hans": language = "zh-Hans" else: language = "en-US" - account = AccountService.get_account_by_email_with_case_fallback(args.email, session=db.session()) + account = AccountService.get_account_by_email_with_case_fallback(req_data.email, session=db.session()) token = AccountService.send_reset_password_email( account=account, @@ -106,16 +106,16 @@ class ForgotPasswordCheckApi(Resource): @console_ns.response(400, "Invalid code or token") @setup_required @email_password_login_enabled - def post(self): - args = ForgotPasswordCheckPayload.model_validate(console_ns.payload) + @model_validate(ForgotPasswordCheckPayload) + def post(self, req_data: ForgotPasswordCheckPayload): - user_email = args.email.lower() + user_email = req_data.email.lower() is_forgot_password_error_rate_limit = AccountService.is_forgot_password_error_rate_limit(user_email) if is_forgot_password_error_rate_limit: raise EmailPasswordResetLimitError() - token_data = AccountService.get_reset_password_data(args.token) + token_data = AccountService.get_reset_password_data(req_data.token) if token_data is None: raise InvalidTokenError() @@ -127,16 +127,16 @@ class ForgotPasswordCheckApi(Resource): if user_email != normalized_token_email: raise InvalidEmailError() - if args.code != token_data.get("code"): + if req_data.code != token_data.get("code"): AccountService.add_forgot_password_error_rate_limit(user_email) raise EmailCodeError() # Verified, revoke the first token - AccountService.revoke_reset_password_token(args.token) + AccountService.revoke_reset_password_token(req_data.token) # Refresh token data by generating a new token _, new_token = AccountService.generate_reset_password_token( - token_email, code=args.code, additional_data={"phase": "reset"} + token_email, code=req_data.code, additional_data={"phase": "reset"} ) AccountService.reset_forgot_password_error_rate_limit(user_email) @@ -156,15 +156,15 @@ class ForgotPasswordResetApi(Resource): @console_ns.response(400, "Invalid token or password mismatch") @setup_required @email_password_login_enabled - def post(self): - args = ForgotPasswordResetPayload.model_validate(console_ns.payload) + @model_validate(ForgotPasswordResetPayload) + def post(self, req_data: ForgotPasswordResetPayload): # Validate passwords match - if args.new_password != args.password_confirm: + if req_data.new_password != req_data.password_confirm: raise PasswordMismatchError() # Validate token and get reset data - reset_data = AccountService.get_reset_password_data(args.token) + reset_data = AccountService.get_reset_password_data(req_data.token) if not reset_data: raise InvalidTokenError() # Must use token in reset phase @@ -172,11 +172,11 @@ class ForgotPasswordResetApi(Resource): raise InvalidTokenError() # Revoke token to prevent reuse - AccountService.revoke_reset_password_token(args.token) + AccountService.revoke_reset_password_token(req_data.token) # Generate secure salt and hash password salt = secrets.token_bytes(16) - password_hashed = hash_password(args.new_password, salt) + password_hashed = hash_password(req_data.new_password, salt) email = reset_data.get("email", "") account = AccountService.get_account_by_email_with_case_fallback(email, session=db.session()) diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 49b248a1e48..9363bc849eb 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -39,6 +39,7 @@ from controllers.console.wraps import ( decrypt_code_field, decrypt_password_field, email_password_login_enabled, + model_validate, setup_required, with_current_user, ) @@ -114,10 +115,10 @@ class LoginApi(Resource): @console_ns.expect(console_ns.models[LoginPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultOptionalDataResponse.__name__]) @decrypt_password_field - def post(self): + @model_validate(LoginPayload) + def post(self, req_data: LoginPayload): """Authenticate user and login.""" - args = LoginPayload.model_validate(console_ns.payload) - request_email = args.email + request_email = req_data.email normalized_email = request_email.lower() if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(normalized_email): @@ -129,7 +130,7 @@ class LoginApi(Resource): _log_console_login_failure(email=normalized_email, reason=LoginFailureReason.LOGIN_RATE_LIMITED) raise EmailPasswordLoginLimitError() - invite_token = args.invite_token + invite_token = req_data.invite_token invitation_data: InvitationDetailDict | None = None if invite_token: invitation_data = RegisterService.get_invitation_with_case_fallback( @@ -150,7 +151,7 @@ class LoginApi(Resource): ) raise InvalidEmailError() account = _authenticate_account_with_case_fallback( - request_email, normalized_email, args.password, invite_token + request_email, normalized_email, req_data.password, invite_token ) except services.errors.account.AccountLoginError: _log_console_login_failure(email=normalized_email, reason=LoginFailureReason.ACCOUNT_BANNED) @@ -215,16 +216,16 @@ class ResetPasswordSendEmailApi(Resource): @email_password_login_enabled @console_ns.expect(console_ns.models[EmailPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultDataResponse.__name__]) - def post(self): - args = EmailPayload.model_validate(console_ns.payload) - normalized_email = args.email.lower() + @model_validate(EmailPayload) + def post(self, req_data: EmailPayload): + normalized_email = req_data.email.lower() - if args.language is not None and args.language == "zh-Hans": + if req_data.language is not None and req_data.language == "zh-Hans": language = "zh-Hans" else: language = "en-US" try: - account = _get_account_with_case_fallback(args.email) + account = _get_account_with_case_fallback(req_data.email) except AccountRegisterError: raise AccountInFreezeError() @@ -243,20 +244,20 @@ class EmailCodeLoginSendEmailApi(Resource): @setup_required @console_ns.expect(console_ns.models[EmailPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultDataResponse.__name__]) - def post(self): - args = EmailPayload.model_validate(console_ns.payload) - normalized_email = args.email.lower() + @model_validate(EmailPayload) + def post(self, req_data: EmailPayload): + normalized_email = req_data.email.lower() ip_address = extract_remote_ip(request) if AccountService.is_email_send_ip_limit(ip_address): raise EmailSendIpLimitError() - if args.language is not None and args.language == "zh-Hans": + if req_data.language is not None and req_data.language == "zh-Hans": language = "zh-Hans" else: language = "en-US" try: - account = _get_account_with_case_fallback(args.email) + account = _get_account_with_case_fallback(req_data.email) except AccountRegisterError: raise AccountInFreezeError() @@ -277,14 +278,14 @@ class EmailCodeLoginApi(Resource): @console_ns.expect(console_ns.models[EmailCodeLoginPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) @decrypt_code_field - def post(self): - args = EmailCodeLoginPayload.model_validate(console_ns.payload) + @model_validate(EmailCodeLoginPayload) + def post(self, req_data: EmailCodeLoginPayload): - original_email = args.email + original_email = req_data.email user_email = original_email.lower() - language = args.language + language = req_data.language - token_data = AccountService.get_email_code_login_data(args.token) + token_data = AccountService.get_email_code_login_data(req_data.token) if token_data is None: _log_console_login_failure(email=user_email, reason=LoginFailureReason.INVALID_EMAIL_CODE_TOKEN) raise InvalidTokenError() @@ -295,11 +296,11 @@ class EmailCodeLoginApi(Resource): _log_console_login_failure(email=user_email, reason=LoginFailureReason.EMAIL_CODE_EMAIL_MISMATCH) raise InvalidEmailError() - if token_data["code"] != args.code: + if token_data["code"] != req_data.code: _log_console_login_failure(email=user_email, reason=LoginFailureReason.INVALID_EMAIL_CODE) raise EmailCodeError() - AccountService.revoke_email_code_login_token(args.token) + AccountService.revoke_email_code_login_token(req_data.token) try: account = _get_account_with_case_fallback(original_email) except Unauthorized as exc: @@ -325,7 +326,7 @@ class EmailCodeLoginApi(Resource): email=user_email, name=user_email, interface_language=get_valid_language(language), - timezone=args.timezone, + timezone=req_data.timezone, session=db.session(), ) except WorkSpaceNotAllowedCreateError: diff --git a/api/controllers/console/billing/billing.py b/api/controllers/console/billing/billing.py index 3a983b50176..aa69f48dbe4 100644 --- a/api/controllers/console/billing/billing.py +++ b/api/controllers/console/billing/billing.py @@ -1,7 +1,6 @@ import base64 from typing import Any, Literal -from flask import request from flask_restx import Resource from pydantic import BaseModel, Field, RootModel from werkzeug.exceptions import BadRequest @@ -10,6 +9,7 @@ from controllers.common.schema import query_params_from_model, register_response from controllers.console import console_ns from controllers.console.wraps import ( account_initialization_required, + model_validate, only_edition_cloud, setup_required, with_current_tenant_id, @@ -54,10 +54,10 @@ class Subscription(Resource): @only_edition_cloud @with_current_user @with_current_tenant_id - def get(self, current_tenant_id: str, current_user: Account): - args = SubscriptionQuery.model_validate(request.args.to_dict(flat=True)) + @model_validate(SubscriptionQuery) + def get(self, req_data: SubscriptionQuery, current_tenant_id: str, current_user: Account): BillingService.is_tenant_owner_or_admin(current_user, session=db.session()) - return BillingService.get_subscription(args.plan, args.interval, current_user.email, current_tenant_id) + return BillingService.get_subscription(req_data.plan, req_data.interval, current_user.email, current_tenant_id) @console_ns.route("/billing/invoices") @@ -87,10 +87,10 @@ class PartnerTenants(Resource): @account_initialization_required @only_edition_cloud @with_current_user - def put(self, current_user: Account, partner_key: str): + @model_validate(PartnerTenantsPayload) + def put(self, req_data: PartnerTenantsPayload, current_user: Account, partner_key: str): try: - args = PartnerTenantsPayload.model_validate(console_ns.payload or {}) - click_id = args.click_id + click_id = req_data.click_id decoded_partner_key = base64.b64decode(partner_key).decode("utf-8") except Exception: raise BadRequest("Invalid partner_key") diff --git a/api/controllers/console/billing/compliance.py b/api/controllers/console/billing/compliance.py index ea5852586a9..f62d6f8e417 100644 --- a/api/controllers/console/billing/compliance.py +++ b/api/controllers/console/billing/compliance.py @@ -14,6 +14,7 @@ from ...common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0 from .. import console_ns from ..wraps import ( account_initialization_required, + model_validate, only_edition_cloud, setup_required, with_current_tenant_id, @@ -48,13 +49,13 @@ class ComplianceApi(Resource): @only_edition_cloud @with_current_user @with_current_tenant_id - def get(self, current_tenant_id: str, current_user: Account): - args = ComplianceDownloadQuery.model_validate(request.args.to_dict(flat=True)) + @model_validate(ComplianceDownloadQuery) + def get(self, req_data: ComplianceDownloadQuery, current_tenant_id: str, current_user: Account): ip_address = extract_remote_ip(request) device_info = request.headers.get("User-Agent", "Unknown device") return BillingService.get_compliance_download_link( - doc_name=args.doc_name, + doc_name=req_data.doc_name, account_id=current_user.id, tenant_id=current_tenant_id, ip=ip_address, diff --git a/api/controllers/console/remote_files.py b/api/controllers/console/remote_files.py index c771f4489ac..e3aefb72595 100644 --- a/api/controllers/console/remote_files.py +++ b/api/controllers/console/remote_files.py @@ -51,8 +51,8 @@ def upload_remote_file_from_request( current_user: Account, resource_tenant_id: str | None = None, ) -> FileWithSignedUrl: - """Validate the JSON request, fetch its remote file, and persist it under the requested tenant.""" payload = RemoteFileUploadPayload.model_validate(console_ns.payload) + """Validate the JSON request, fetch its remote file, and persist it under the requested tenant.""" url = payload.url # Try to fetch remote file metadata/content first diff --git a/api/controllers/console/tag/tags.py b/api/controllers/console/tag/tags.py index 00596545d1c..0084e614c5b 100644 --- a/api/controllers/console/tag/tags.py +++ b/api/controllers/console/tag/tags.py @@ -1,7 +1,6 @@ from typing import Literal from uuid import UUID -from flask import request from flask_restx import Resource from pydantic import BaseModel, Field, RootModel, field_validator from sqlalchemy import select @@ -135,10 +134,9 @@ class TagListApi(Resource): @console_ns.doc(params=query_params_from_model(TagListQueryParam)) @console_ns.response(200, "Success", console_ns.models[TagListResponse.__name__]) @with_current_tenant_id - def get(self, current_tenant_id: str): - raw_args = request.args.to_dict() - param = TagListQueryParam.model_validate(raw_args) - tags = TagService.get_tags(param.type, current_tenant_id, param.keyword, session=db.session()) + @model_validate(TagListQueryParam) + def get(self, req_data: TagListQueryParam, current_tenant_id: str): + tags = TagService.get_tags(req_data.type, current_tenant_id, req_data.keyword, session=db.session()) return dump_response(TagListResponse, tags), 200 @@ -213,10 +211,9 @@ def _require_tag_binding_edit_permission(current_user: Account) -> None: raise Forbidden() -def _create_tag_bindings(current_user: Account) -> tuple[dict[str, str], int]: +def _create_tag_bindings(current_user: Account, payload: TagBindingPayload) -> tuple[dict[str, str], int]: _require_tag_binding_edit_permission(current_user) - payload = TagBindingPayload.model_validate(console_ns.payload or {}) _enforce_snippet_tag_rbac_if_needed(payload.type) TagService.save_tag_binding( TagBindingCreatePayload( @@ -229,10 +226,9 @@ def _create_tag_bindings(current_user: Account) -> tuple[dict[str, str], int]: return {"result": "success"}, 200 -def _remove_tag_bindings(current_user: Account) -> tuple[dict[str, str], int]: +def _remove_tag_bindings(current_user: Account, payload: TagBindingRemovePayload) -> tuple[dict[str, str], int]: _require_tag_binding_edit_permission(current_user) - payload = TagBindingRemovePayload.model_validate(console_ns.payload or {}) _enforce_snippet_tag_rbac_if_needed(payload.type) TagService.delete_tag_binding( TagBindingDeletePayload( @@ -256,8 +252,9 @@ class TagBindingCollectionApi(Resource): @login_required @account_initialization_required @with_current_user - def post(self, current_user: Account): - return _create_tag_bindings(current_user) + @model_validate(TagBindingPayload) + def post(self, req_data: TagBindingPayload, current_user: Account): + return _create_tag_bindings(current_user, req_data) @console_ns.route("/tag-bindings/remove") @@ -272,5 +269,6 @@ class TagBindingRemoveApi(Resource): @login_required @account_initialization_required @with_current_user - def post(self, current_user: Account): - return _remove_tag_bindings(current_user) + @model_validate(TagBindingRemovePayload) + def post(self, req_data: TagBindingRemovePayload, current_user: Account): + return _remove_tag_bindings(current_user, req_data) diff --git a/api/tests/unit_tests/controllers/console/billing/test_billing.py b/api/tests/unit_tests/controllers/console/billing/test_billing.py index 6521f4d852c..87f90caa800 100644 --- a/api/tests/unit_tests/controllers/console/billing/test_billing.py +++ b/api/tests/unit_tests/controllers/console/billing/test_billing.py @@ -5,7 +5,7 @@ from unittest.mock import patch import pytest from flask import Flask from sqlalchemy.orm import Session -from werkzeug.exceptions import BadRequest +from werkzeug.exceptions import BadRequest, UnprocessableEntity from controllers.console import wraps as console_wraps from controllers.console.billing.billing import PartnerTenants @@ -129,7 +129,7 @@ class TestPartnerTenants: assert "Invalid partner_key" in str(exc_info.value) def test_put_missing_click_id(self, app: Flask, mock_account, mock_billing_service, mock_decorators): - """Test that missing click_id raises BadRequest.""" + """Test that missing click_id raises UnprocessableEntity (422).""" # Arrange partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8") @@ -148,8 +148,8 @@ class TestPartnerTenants: resource = PartnerTenants() # Act & Assert - # Validation should raise BadRequest for missing required field - with pytest.raises(BadRequest): + # Validation should raise UnprocessableEntity (422) for missing required field + with pytest.raises(UnprocessableEntity): resource.put(partner_key_encoded) def test_put_billing_service_json_decode_error( 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 55adfab959f..aa04caed752 100644 --- a/api/tests/unit_tests/controllers/console/tag/test_tags.py +++ b/api/tests/unit_tests/controllers/console/tag/test_tags.py @@ -13,8 +13,11 @@ from controllers.console import console_ns from controllers.console.tag.tags import ( TagBasePayload, TagBindingCollectionApi, + TagBindingPayload, TagBindingRemoveApi, + TagBindingRemovePayload, TagListApi, + TagListQueryParam, TagUpdateDeleteApi, TagUpdateRequestPayload, ) @@ -132,7 +135,7 @@ class TestTagListApi: ], ), ): - result, status = method(api, "tenant-1") + result, status = method(api, TagListQueryParam(type="knowledge"), "tenant-1") assert status == 200 assert result == [{"id": "1", "name": "tag", "type": "knowledge", "binding_count": "1"}] @@ -155,7 +158,7 @@ class TestTagListApi: ], ) as get_tags_mock, ): - result, status = method(api, "tenant-1") + result, status = method(api, TagListQueryParam(type="snippet"), "tenant-1") get_tags_mock.assert_called_once() assert get_tags_mock.call_args.args == ("snippet", "tenant-1", None) @@ -219,7 +222,7 @@ class TestTagListApi: with app.test_request_context("/"): with pytest.raises(Forbidden): - method(api, None, readonly_user) + method(api, TagBasePayload(name="test", type=TagType.KNOWLEDGE), readonly_user) class TestTagUpdateDeleteApi: @@ -256,7 +259,7 @@ class TestTagUpdateDeleteApi: with app.test_request_context("/"): with pytest.raises(Forbidden): - method(api, None, readonly_user, "tag-1") + method(api, TagUpdateRequestPayload(name="test"), readonly_user, "tag-1") def test_delete_success(self, app: Flask, admin_user, sqlite_engine: Engine): api = TagUpdateDeleteApi() @@ -375,7 +378,7 @@ class TestTagBindingCollectionApi: payload_patch(payload), patch("controllers.console.tag.tags.TagService.save_tag_binding") as save_mock, ): - result, status = method(api, admin_user) + result, status = method(api, TagBindingPayload.model_validate(payload), admin_user) save_mock.assert_called_once() assert status == 200 @@ -396,7 +399,7 @@ class TestTagBindingCollectionApi: payload_patch(payload), patch("controllers.console.tag.tags.TagService.save_tag_binding") as save_mock, ): - result, status = method(api, admin_user) + result, status = method(api, TagBindingPayload.model_validate(payload), admin_user) save_mock.assert_called_once() binding_payload = save_mock.call_args.args[0] @@ -414,7 +417,11 @@ class TestTagBindingCollectionApi: payload_patch({}), ): with pytest.raises(Forbidden): - method(api, readonly_user) + method( + api, + TagBindingPayload(tag_ids=["tag-1"], target_id="target-1", type=TagType.KNOWLEDGE), + readonly_user, + ) class TestTagBindingRemoveApi: @@ -433,7 +440,7 @@ class TestTagBindingRemoveApi: payload_patch(payload), patch("controllers.console.tag.tags.TagService.delete_tag_binding") as delete_mock, ): - result, status = method(api, admin_user) + result, status = method(api, TagBindingRemovePayload.model_validate(payload), admin_user) delete_mock.assert_called_once() delete_payload = delete_mock.call_args.args[0] @@ -450,7 +457,11 @@ class TestTagBindingRemoveApi: payload_patch({}), ): with pytest.raises(Forbidden): - method(api, readonly_user) + method( + api, + TagBindingRemovePayload(tag_ids=["tag-1"], target_id="target-1", type=TagType.KNOWLEDGE), + readonly_user, + ) class TestTagResponseModel: