diff --git a/api/controllers/web/saved_message.py b/api/controllers/web/saved_message.py index e085b6ece1a..eed4aed5ed6 100644 --- a/api/controllers/web/saved_message.py +++ b/api/controllers/web/saved_message.py @@ -1,15 +1,16 @@ from uuid import UUID from pydantic import TypeAdapter +from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound from controllers.common.controller_schemas import SavedMessageCreatePayload, SavedMessageListQuery from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models +from controllers.console.app.wraps import with_session from controllers.console.wraps import model_validate from controllers.web import web_ns from controllers.web.error import NotCompletionAppError from controllers.web.wraps import WebApiResource -from extensions.ext_database import db from fields.conversation_fields import MessageResponseSource, ResultResponse from fields.message_fields import SavedMessageInfiniteScrollPagination, SavedMessageItem from models.model import App, EndUser @@ -36,12 +37,12 @@ class SavedMessageListApi(WebApiResource): } ) @web_ns.response(200, "Success", web_ns.models[SavedMessageInfiniteScrollPagination.__name__]) + @with_session(write=False) @model_validate(SavedMessageListQuery) - def get(self, query: SavedMessageListQuery, app_model: App, end_user: EndUser): + def get(self, query: SavedMessageListQuery, session: Session, app_model: App, end_user: EndUser): if app_model.mode != "completion": raise NotCompletionAppError() - session = db.session() pagination = SavedMessageService.pagination_by_last_id( app_model, end_user, query.last_id, query.limit, session=session ) @@ -73,13 +74,14 @@ class SavedMessageListApi(WebApiResource): ) @web_ns.response(200, "Message saved successfully", web_ns.models[ResultResponse.__name__]) @web_ns.expect(web_ns.models[SavedMessageCreatePayload.__name__]) + @with_session @model_validate(SavedMessageCreatePayload) - def post(self, payload: SavedMessageCreatePayload, app_model: App, end_user: EndUser): + def post(self, payload: SavedMessageCreatePayload, session: Session, app_model: App, end_user: EndUser): if app_model.mode != "completion": raise NotCompletionAppError() try: - SavedMessageService.save(app_model, end_user, payload.message_id, session=db.session()) + SavedMessageService.save(app_model, end_user, payload.message_id, session=session) except MessageNotExistsError: raise NotFound("Message Not Exists.") @@ -102,12 +104,13 @@ class SavedMessageApi(WebApiResource): } ) @web_ns.response(204, "Message removed successfully") - def delete(self, app_model: App, end_user: EndUser, message_id: UUID): + @with_session + def delete(self, session: Session, app_model: App, end_user: EndUser, message_id: UUID): message_id_str = str(message_id) if app_model.mode != "completion": raise NotCompletionAppError() - SavedMessageService.delete(app_model, end_user, message_id_str, session=db.session()) + SavedMessageService.delete(app_model, end_user, message_id_str, session=session) return "", 204 diff --git a/api/tests/unit_tests/controllers/web/test_saved_message.py b/api/tests/unit_tests/controllers/web/test_saved_message.py index fdb276f8f5a..22ec43a5b47 100644 --- a/api/tests/unit_tests/controllers/web/test_saved_message.py +++ b/api/tests/unit_tests/controllers/web/test_saved_message.py @@ -36,9 +36,9 @@ def _end_user() -> EndUser: ) -# The @model_validate decorator wraps the handler; tests call the undecorated -# function so they can pass the validated pydantic payload directly and skip -# the flask request-parsing step. +# The @with_session and @model_validate decorators wrap the handler; tests call +# the undecorated function so they can pass the validated pydantic payload and a +# stub session directly and skip the flask request-parsing step. _list_get = inspect.unwrap(SavedMessageListApi.get) _list_post = inspect.unwrap(SavedMessageListApi.post) @@ -51,7 +51,7 @@ class TestSavedMessageListApiGet: query = SavedMessageListQuery.model_validate({}) with app.test_request_context("/saved-messages"): with pytest.raises(NotCompletionAppError): - _list_get(SavedMessageListApi(), query, _chat_app(), _end_user()) + _list_get(SavedMessageListApi(), query, MagicMock(), _chat_app(), _end_user()) @patch("controllers.web.saved_message.SavedMessageService.pagination_by_last_id") def test_happy_path(self, mock_paginate: MagicMock, app: Flask) -> None: @@ -59,7 +59,7 @@ class TestSavedMessageListApiGet: query = SavedMessageListQuery.model_validate({"limit": 20}) with app.test_request_context("/saved-messages?limit=20"): - result = _list_get(SavedMessageListApi(), query, _completion_app(), _end_user()) + result = _list_get(SavedMessageListApi(), query, MagicMock(), _completion_app(), _end_user()) assert result["limit"] == 20 assert result["has_more"] is False @@ -73,14 +73,14 @@ class TestSavedMessageListApiPost: payload = SavedMessageCreatePayload.model_validate({"message_id": str(uuid4())}) with app.test_request_context("/saved-messages", method="POST"): with pytest.raises(NotCompletionAppError): - _list_post(SavedMessageListApi(), payload, _chat_app(), _end_user()) + _list_post(SavedMessageListApi(), payload, MagicMock(), _chat_app(), _end_user()) @patch("controllers.web.saved_message.SavedMessageService.save") def test_save_success(self, mock_save: MagicMock, app: Flask) -> None: payload = SavedMessageCreatePayload.model_validate({"message_id": str(uuid4())}) with app.test_request_context("/saved-messages", method="POST"): - result = _list_post(SavedMessageListApi(), payload, _completion_app(), _end_user()) + result = _list_post(SavedMessageListApi(), payload, MagicMock(), _completion_app(), _end_user()) assert result["result"] == "success" @@ -90,7 +90,7 @@ class TestSavedMessageListApiPost: with app.test_request_context("/saved-messages", method="POST"): with pytest.raises(NotFound, match="Message Not Exists"): - _list_post(SavedMessageListApi(), payload, _completion_app(), _end_user()) + _list_post(SavedMessageListApi(), payload, MagicMock(), _completion_app(), _end_user()) # ---------------------------------------------------------------------------