From 3a9636faaddd8a2f5d375fa7253e60bf9d892440 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 19 Sep 2026 09:26:43 +0000 Subject: [PATCH] test: migrate console snippets sessions and ORM models to SQLite (#40565) --- .../console/snippets/snippet_workflow.py | 141 ++++----- .../snippet_workflow_draft_variable.py | 97 ++++--- api/controllers/console/workspace/snippets.py | 90 +++--- .../console/snippets/test_snippet_workflow.py | 274 ++++++------------ .../test_snippet_workflow_draft_variable.py | 180 +++++------- .../console/workspace/test_snippets.py | 61 ++-- 6 files changed, 361 insertions(+), 482 deletions(-) diff --git a/api/controllers/console/snippets/snippet_workflow.py b/api/controllers/console/snippets/snippet_workflow.py index d4a6497ad7a..076a6945ff4 100644 --- a/api/controllers/console/snippets/snippet_workflow.py +++ b/api/controllers/console/snippets/snippet_workflow.py @@ -5,7 +5,7 @@ from functools import wraps from flask import request from flask_restx import Resource from pydantic import BaseModel, Field -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from werkzeug.exceptions import BadRequest, InternalServerError, NotFound from controllers.common.controller_schemas import WorkflowUpdatePayload @@ -45,7 +45,7 @@ from controllers.console.wraps import ( ) from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.entities.app_invoke_entities import InvokeFrom -from extensions.ext_database import db +from core.db.session_factory import session_factory from extensions.ext_redis import redis_client from fields.workflow_run_fields import ( WorkflowRunDetailResponse, @@ -73,12 +73,8 @@ logger = logging.getLogger(__name__) # Register Pydantic models with Swagger -def _snippet_session_maker() -> sessionmaker[Session]: - return sessionmaker(bind=db.engine, expire_on_commit=False) - - def _snippet_service() -> SnippetService: - return SnippetService(_snippet_session_maker()) + return SnippetService(session_factory.get_session_maker()) class SnippetWorkflowResponse(WorkflowResponse): @@ -173,9 +169,10 @@ class SnippetDraftWorkflowApi(Resource): @setup_required @login_required @account_initialization_required + @with_session(write=False) @get_snippet @edit_permission_required - def get(self, snippet: CustomizedSnippet): + def get(self, session: Session, snippet: CustomizedSnippet): """Get draft workflow for snippet.""" snippet_service = _snippet_service() workflow = snippet_service.get_draft_workflow(snippet=snippet) @@ -184,7 +181,6 @@ class SnippetDraftWorkflowApi(Resource): raise DraftWorkflowNotExist() workflow.conversation_variables = [] - session = db.session() response = SnippetWorkflowResponse.model_validate( WorkflowResponseSource(workflow, session=session), from_attributes=True ).model_dump(mode="json") @@ -267,9 +263,10 @@ class SnippetPublishedWorkflowApi(Resource): @setup_required @login_required @account_initialization_required + @with_session(write=False) @get_snippet @edit_permission_required - def get(self, snippet: CustomizedSnippet): + def get(self, session: Session, snippet: CustomizedSnippet): """Get published workflow for snippet.""" if not snippet.is_published: return None @@ -281,7 +278,7 @@ class SnippetPublishedWorkflowApi(Resource): return None response = SnippetWorkflowResponse.model_validate( - WorkflowResponseSource(workflow, session=db.session()), from_attributes=True + WorkflowResponseSource(workflow, session=session), from_attributes=True ).model_dump(mode="json") response["input_fields"] = snippet.input_fields_list return response @@ -294,25 +291,27 @@ class SnippetPublishedWorkflowApi(Resource): @login_required @account_initialization_required @with_current_user + @with_session @get_snippet @edit_permission_required @rbac_permission_required(RBACCheck(RBACPermission.SNIPPETS_CREATE_AND_MODIFY, Workspace())) - def post(self, current_user: Account, snippet: CustomizedSnippet): + def post(self, session: Session, current_user: Account, snippet: CustomizedSnippet): """Publish snippet workflow.""" snippet_service = _snippet_service() - with Session(db.engine) as session: - snippet = session.merge(snippet) - try: - workflow = snippet_service.publish_workflow( - session=session, - snippet=snippet, - account=current_user, - ) - workflow_created_at = TimestampField().format(workflow.created_at) - session.commit() - except ValueError as e: - return {"message": str(e)}, 400 + snippet = session.merge(snippet) + try: + workflow = snippet_service.publish_workflow( + session=session, + snippet=snippet, + account=current_user, + ) + workflow_created_at = TimestampField().format(workflow.created_at) + except ValueError as e: + # This domain validation remains a 400 response, so it cannot escape to + # the request-session decorator that normally owns rollback-on-error. + session.rollback() # guard-ignore: no-new-controller-sqlalchemy -- translated validation response + return {"message": str(e)}, 400 return { "result": "success", @@ -353,31 +352,36 @@ class SnippetPublishedAllWorkflowApi(Resource): @setup_required @login_required @account_initialization_required + @with_session(write=False) @get_snippet @edit_permission_required @rbac_permission_required(RBACCheck(RBACPermission.SNIPPETS_CREATE_AND_MODIFY, Workspace())) @model_validate(SnippetWorkflowListQuery) - def get(self, req_data: SnippetWorkflowListQuery, snippet: CustomizedSnippet): + def get( + self, + req_data: SnippetWorkflowListQuery, + session: Session, + snippet: CustomizedSnippet, + ): """Get all published workflow versions for snippet.""" snippet_service = _snippet_service() - with Session(db.engine) as session: - workflows, has_more = snippet_service.get_all_published_workflows( - session=session, - snippet=snippet, - page=req_data.page, - limit=req_data.limit, - ) + workflows, has_more = snippet_service.get_all_published_workflows( + session=session, + snippet=snippet, + page=req_data.page, + limit=req_data.limit, + ) - response = SnippetWorkflowPaginationResponse.model_validate( - { - "items": [WorkflowResponseSource(workflow, session=session) for workflow in workflows], - "page": req_data.page, - "limit": req_data.limit, - "has_more": has_more, - }, - from_attributes=True, - ).model_dump(mode="json") + response = SnippetWorkflowPaginationResponse.model_validate( + { + "items": [WorkflowResponseSource(workflow, session=session) for workflow in workflows], + "page": req_data.page, + "limit": req_data.limit, + "has_more": has_more, + }, + from_attributes=True, + ).model_dump(mode="json") for item in response["items"]: item["input_fields"] = snippet.input_fields_list return response @@ -435,6 +439,7 @@ class SnippetWorkflowByIdApi(Resource): @login_required @account_initialization_required @with_current_user + @with_session @get_snippet @edit_permission_required @rbac_permission_required(RBACCheck(RBACPermission.SNIPPETS_CREATE_AND_MODIFY, Workspace())) @@ -442,6 +447,7 @@ class SnippetWorkflowByIdApi(Resource): def patch( self, req_data: WorkflowUpdatePayload, + session: Session, current_user: Account, snippet: CustomizedSnippet, workflow_id: str, @@ -453,22 +459,21 @@ class SnippetWorkflowByIdApi(Resource): return {"message": "No valid fields to update"}, 400 snippet_service = _snippet_service() - with _snippet_session_maker().begin() as session: - workflow = snippet_service.update_workflow( - session=session, - snippet=snippet, - workflow_id=workflow_id, - account=current_user, - data=update_data, - ) - if not workflow: - raise NotFound("Workflow not found") + workflow = snippet_service.update_workflow( + session=session, + snippet=snippet, + workflow_id=workflow_id, + account=current_user, + data=update_data, + ) + if not workflow: + raise NotFound("Workflow not found") - response = SnippetWorkflowResponse.model_validate( - WorkflowResponseSource(workflow, session=session), from_attributes=True - ).model_dump(mode="json") - response["input_fields"] = snippet.input_fields_list - return response + response = SnippetWorkflowResponse.model_validate( + WorkflowResponseSource(workflow, session=session), from_attributes=True + ).model_dump(mode="json") + response["input_fields"] = snippet.input_fields_list + return response @console_ns.doc("delete_snippet_workflow_by_id") @console_ns.doc(description="Delete a published snippet workflow version") @@ -485,7 +490,7 @@ class SnippetWorkflowByIdApi(Resource): def delete(self, snippet: CustomizedSnippet, workflow_id: str): """Delete a published snippet workflow version.""" snippet_service = _snippet_service() - with _snippet_session_maker().begin() as session: + with session_factory.get_session_maker().begin() as session: try: snippet_service.delete_workflow( session=session, @@ -641,12 +646,13 @@ class SnippetDraftNodeRunApi(Resource): account=current_user, query=req_data.query, files=files, - session_maker=_snippet_session_maker(), + session_maker=session_factory.get_session_maker(), ) - return WorkflowRunNodeExecutionResponse.model_validate( - node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True - ).model_dump(mode="json") + with session_factory.create_session() as session: + return WorkflowRunNodeExecutionResponse.model_validate( + node_execution_response_source(workflow_node_execution, session=session), from_attributes=True + ).model_dump(mode="json") @console_ns.route("/snippets//workflows/draft/nodes//last-run") @@ -682,9 +688,10 @@ class SnippetDraftNodeLastRunApi(Resource): if node_exec is None: raise NotFound("Node last run not found") - return WorkflowRunNodeExecutionResponse.model_validate( - node_execution_response_source(node_exec, session=db.session()), from_attributes=True - ).model_dump(mode="json") + with session_factory.create_session() as session: + return WorkflowRunNodeExecutionResponse.model_validate( + node_execution_response_source(node_exec, session=session), from_attributes=True + ).model_dump(mode="json") @console_ns.route("/snippets//workflows/draft/iteration/nodes//run") @@ -728,7 +735,7 @@ class SnippetDraftRunIterationNodeApi(Resource): node_id=node_id, args=args, streaming=True, - session_maker=_snippet_session_maker(), + session_maker=session_factory.get_session_maker(), ) return helper.compact_generate_response(response) @@ -779,7 +786,7 @@ class SnippetDraftRunLoopNodeApi(Resource): node_id=node_id, args=req_data, streaming=True, - session_maker=_snippet_session_maker(), + session_maker=session_factory.get_session_maker(), ) return helper.compact_generate_response(response) @@ -823,7 +830,7 @@ class SnippetDraftWorkflowRunApi(Resource): args=args, invoke_from=InvokeFrom.DEBUGGER, streaming=True, - session_maker=_snippet_session_maker(), + session_maker=session_factory.get_session_maker(), ) return helper.compact_generate_response(response) diff --git a/api/controllers/console/snippets/snippet_workflow_draft_variable.py b/api/controllers/console/snippets/snippet_workflow_draft_variable.py index 7091cd24c24..f3cf9c0a284 100644 --- a/api/controllers/console/snippets/snippet_workflow_draft_variable.py +++ b/api/controllers/console/snippets/snippet_workflow_draft_variable.py @@ -16,10 +16,11 @@ from typing import Any, Concatenate from flask import Response from flask_restx import Resource -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.schema import query_params_from_model +from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.app.error import DraftWorkflowNotExist from controllers.console.app.workflow_draft_variable import ( @@ -38,9 +39,9 @@ from controllers.console.wraps import ( with_current_user, ) from core.app.file_access import DatabaseFileAccessController +from core.db.session_factory import session_factory from core.workflow.llm_environment_variable import environment_variable_value_type from core.workflow.variable_prefixes import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID -from extensions.ext_database import db from factories.file_factory import build_from_mapping, build_from_mappings from factories.variable_factory import build_segment_with_type from fields.workflow_draft_variable_fields import ( @@ -64,7 +65,7 @@ _file_access_controller = DatabaseFileAccessController() def _snippet_service() -> SnippetService: - return SnippetService(sessionmaker(bind=db.engine, expire_on_commit=False)) + return SnippetService(session_factory.get_session_maker()) def _ensure_snippet_draft_variable_row_allowed( @@ -107,22 +108,26 @@ class SnippetWorkflowVariableCollectionApi(Resource): ) @_snippet_draft_var_prerequisite @model_validate(WorkflowDraftVariableListQuery) + @with_session(write=False) def get( - self, args: WorkflowDraftVariableListQuery, current_user: Account, snippet: CustomizedSnippet + self, + session: Session, + args: WorkflowDraftVariableListQuery, + current_user: Account, + snippet: CustomizedSnippet, ) -> dict[str, Any]: snippet_service = _snippet_service() if snippet_service.get_draft_workflow(snippet=snippet) is None: raise DraftWorkflowNotExist() - with Session(bind=db.engine, expire_on_commit=False) as session: - draft_var_srv = WorkflowDraftVariableService(session=session) - workflow_vars = draft_var_srv.list_variables_without_values( - app_id=snippet.id, - page=args.page, - limit=args.limit, - user_id=current_user.id, - exclude_node_ids=_SNIPPET_EXCLUDED_DRAFT_VARIABLE_NODE_IDS, - ) + draft_var_srv = WorkflowDraftVariableService(session=session) + workflow_vars = draft_var_srv.list_variables_without_values( + app_id=snippet.id, + page=args.page, + limit=args.limit, + user_id=current_user.id, + exclude_node_ids=_SNIPPET_EXCLUDED_DRAFT_VARIABLE_NODE_IDS, + ) return dump_response(WorkflowDraftVariableListWithoutValueResponse, workflow_vars) @@ -130,10 +135,10 @@ class SnippetWorkflowVariableCollectionApi(Resource): @console_ns.doc(description="Delete all draft workflow variables for the current user (snippet scope)") @console_ns.response(204, "Workflow variables deleted successfully") @_snippet_draft_var_prerequisite - def delete(self, current_user: Account, snippet: CustomizedSnippet) -> Response: - draft_var_srv = WorkflowDraftVariableService(session=db.session()) + @with_session + def delete(self, session: Session, current_user: Account, snippet: CustomizedSnippet) -> Response: + draft_var_srv = WorkflowDraftVariableService(session=session) draft_var_srv.delete_user_workflow_variables(snippet.id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -147,11 +152,11 @@ class SnippetNodeVariableCollectionApi(Resource): console_ns.models[WorkflowDraftVariableListResponse.__name__], ) @_snippet_draft_var_prerequisite - def get(self, current_user: Account, snippet: CustomizedSnippet, node_id: str) -> dict[str, Any]: + @with_session(write=False) + def get(self, session: Session, current_user: Account, snippet: CustomizedSnippet, node_id: str) -> dict[str, Any]: validate_node_id(node_id) - with Session(bind=db.engine, expire_on_commit=False) as session: - draft_var_srv = WorkflowDraftVariableService(session=session) - node_vars = draft_var_srv.list_node_variables(snippet.id, node_id, user_id=current_user.id) + draft_var_srv = WorkflowDraftVariableService(session=session) + node_vars = draft_var_srv.list_node_variables(snippet.id, node_id, user_id=current_user.id) return dump_response(WorkflowDraftVariableListResponse, node_vars) @@ -159,11 +164,11 @@ class SnippetNodeVariableCollectionApi(Resource): @console_ns.doc(description="Delete all variables for a specific node (snippet draft workflow)") @console_ns.response(204, "Node variables deleted successfully") @_snippet_draft_var_prerequisite - def delete(self, current_user: Account, snippet: CustomizedSnippet, node_id: str) -> Response: + @with_session + def delete(self, session: Session, current_user: Account, snippet: CustomizedSnippet, node_id: str) -> Response: validate_node_id(node_id) - srv = WorkflowDraftVariableService(db.session()) + srv = WorkflowDraftVariableService(session) srv.delete_node_variables(snippet.id, node_id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -178,8 +183,11 @@ class SnippetVariableApi(Resource): ) @console_ns.response(404, "Variable not found") @_snippet_draft_var_prerequisite - def get(self, current_user: Account, snippet: CustomizedSnippet, variable_id: str) -> dict[str, Any]: - draft_var_srv = WorkflowDraftVariableService(session=db.session()) + @with_session(write=False) + def get( + self, session: Session, current_user: Account, snippet: CustomizedSnippet, variable_id: str + ) -> dict[str, Any]: + draft_var_srv = WorkflowDraftVariableService(session=session) variable = ensure_variable_access( variable=draft_var_srv.get_variable(variable_id=variable_id), app_id=snippet.id, @@ -199,15 +207,17 @@ class SnippetVariableApi(Resource): ) @console_ns.response(404, "Variable not found") @_snippet_draft_var_prerequisite + @with_session @model_validate(WorkflowDraftVariableUpdatePayload) def patch( self, req_data: WorkflowDraftVariableUpdatePayload, + session: Session, current_user: Account, snippet: CustomizedSnippet, variable_id: str, ) -> dict[str, Any]: - draft_var_srv = WorkflowDraftVariableService(session=db.session()) + draft_var_srv = WorkflowDraftVariableService(session=session) variable = ensure_variable_access( variable=draft_var_srv.get_variable(variable_id=variable_id), @@ -227,24 +237,23 @@ class SnippetVariableApi(Resource): if variable.value_type == SegmentType.FILE: if not isinstance(raw_value, dict): raise InvalidArgumentError(description=f"expected dict for file, got {type(raw_value)}") - raw_value = build_from_mapping( - mapping=raw_value, - tenant_id=snippet.tenant_id, - access_controller=_file_access_controller, - ) + raw_value = build_from_mapping( + mapping=raw_value, + tenant_id=snippet.tenant_id, + access_controller=_file_access_controller, + ) elif variable.value_type == SegmentType.ARRAY_FILE: if not isinstance(raw_value, list): raise InvalidArgumentError(description=f"expected list for files, got {type(raw_value)}") if len(raw_value) > 0 and not isinstance(raw_value[0], dict): raise InvalidArgumentError(description=f"expected dict for files[0], got {type(raw_value)}") - raw_value = build_from_mappings( - mappings=raw_value, - tenant_id=snippet.tenant_id, - access_controller=_file_access_controller, - ) + raw_value = build_from_mappings( + mappings=raw_value, + tenant_id=snippet.tenant_id, + access_controller=_file_access_controller, + ) new_value = build_segment_with_type(variable.value_type, raw_value) draft_var_srv.update_variable(variable, name=new_name, value=new_value) - db.session.commit() return dump_response(WorkflowDraftVariableResponse, variable) @console_ns.doc("delete_snippet_workflow_variable") @@ -252,8 +261,9 @@ class SnippetVariableApi(Resource): @console_ns.response(204, "Variable deleted successfully") @console_ns.response(404, "Variable not found") @_snippet_draft_var_prerequisite - def delete(self, current_user: Account, snippet: CustomizedSnippet, variable_id: str) -> Response: - draft_var_srv = WorkflowDraftVariableService(session=db.session()) + @with_session + def delete(self, session: Session, current_user: Account, snippet: CustomizedSnippet, variable_id: str) -> Response: + draft_var_srv = WorkflowDraftVariableService(session=session) variable = ensure_variable_access( variable=draft_var_srv.get_variable(variable_id=variable_id), app_id=snippet.id, @@ -262,7 +272,6 @@ class SnippetVariableApi(Resource): ) _ensure_snippet_draft_variable_row_allowed(variable=variable, variable_id=variable_id) draft_var_srv.delete_variable(variable) - db.session.commit() return Response("", 204) @@ -278,8 +287,11 @@ class SnippetVariableResetApi(Resource): @console_ns.response(204, "Variable reset (no content)") @console_ns.response(404, "Variable not found") @_snippet_draft_var_prerequisite - def put(self, current_user: Account, snippet: CustomizedSnippet, variable_id: str) -> Response | Any: - draft_var_srv = WorkflowDraftVariableService(session=db.session()) + @with_session + def put( + self, session: Session, current_user: Account, snippet: CustomizedSnippet, variable_id: str + ) -> Response | Any: + draft_var_srv = WorkflowDraftVariableService(session=session) snippet_service = _snippet_service() draft_workflow = snippet_service.get_draft_workflow(snippet=snippet) if draft_workflow is None: @@ -295,7 +307,6 @@ class SnippetVariableResetApi(Resource): _ensure_snippet_draft_variable_row_allowed(variable=variable, variable_id=variable_id) resetted = draft_var_srv.reset_variable(draft_workflow, variable) - db.session.commit() if resetted is None: return Response("", 204) return dump_response(WorkflowDraftVariableResponse, resetted) diff --git a/api/controllers/console/workspace/snippets.py b/api/controllers/console/workspace/snippets.py index 2b8f609801e..ec69ddf85ea 100644 --- a/api/controllers/console/workspace/snippets.py +++ b/api/controllers/console/workspace/snippets.py @@ -4,7 +4,7 @@ from uuid import UUID from flask import Response, request from flask_restx import Resource -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from werkzeug.exceptions import BadRequest, NotFound from controllers.common.fields import TextFileResponse @@ -33,8 +33,8 @@ from controllers.console.wraps import ( with_current_tenant_id, with_current_user, ) +from core.db.session_factory import session_factory from core.plugin.entities.plugin import PluginDependency -from extensions.ext_database import db from fields.base import ResponseModel from fields.snippet_fields import ( SnippetListItemResponse, @@ -73,8 +73,10 @@ class SnippetUseCountResponse(ResponseModel): use_count: int -def _snippet_service() -> SnippetService: - return SnippetService(sessionmaker(bind=db.engine, expire_on_commit=False)) +def _snippet_service(session: Session | None = None) -> SnippetService: + if session is not None: + return SnippetService(session=session) + return SnippetService(session_factory.get_session_maker()) def _snippet_list_query_from_request() -> SnippetListQuery: @@ -123,7 +125,7 @@ class CustomizedSnippetsApi(Resource): """List customized snippets with pagination and search.""" query = _snippet_list_query_from_request() - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippets, total, has_more = snippet_service.get_snippets( tenant_id=current_tenant_id, session=session, @@ -159,7 +161,13 @@ class CustomizedSnippetsApi(Resource): @with_current_tenant_id @with_session @model_validate(CreateSnippetPayload) - def post(self, req_data: CreateSnippetPayload, session: Session, current_tenant_id: str, current_user: Account): + def post( + self, + req_data: CreateSnippetPayload, + session: Session, + current_tenant_id: str, + current_user: Account, + ): """Create a new customized snippet.""" try: snippet_type = SnippetType(req_data.type) @@ -170,7 +178,7 @@ class CustomizedSnippetsApi(Resource): if req_data.graph is not None: SnippetService.validate_snippet_graph_forbidden_nodes(req_data.graph) - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippet = snippet_service.create_snippet( tenant_id=current_tenant_id, name=req_data.name, @@ -198,7 +206,7 @@ class CustomizedSnippetDetailApi(Resource): @with_session(write=False) def get(self, session: Session, current_tenant_id: str, snippet_id: UUID): """Get customized snippet details.""" - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippet = snippet_service.get_snippet_by_id( snippet_id=str(snippet_id), tenant_id=current_tenant_id, @@ -232,7 +240,7 @@ class CustomizedSnippetDetailApi(Resource): snippet_id: str, ): """Update customized snippet.""" - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippet = snippet_service.get_snippet_by_id( snippet_id=snippet_id, tenant_id=current_tenant_id, @@ -277,9 +285,10 @@ class CustomizedSnippetDetailApi(Resource): @rbac_permission_required(RBACCheck(RBACPermission.SNIPPETS_MANAGE, Workspace())) @with_current_user @with_current_tenant_id - def delete(self, current_tenant_id: str, current_user: Account, snippet_id: str): + @with_session + def delete(self, session: Session, current_tenant_id: str, current_user: Account, snippet_id: str): """Delete customized snippet.""" - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippet = snippet_service.get_snippet_by_id( snippet_id=snippet_id, tenant_id=current_tenant_id, @@ -288,14 +297,11 @@ class CustomizedSnippetDetailApi(Resource): if not snippet: raise NotFound("Snippet not found") - with Session(db.engine) as session: - snippet = session.merge(snippet) - SnippetService.delete_snippet( - session=session, - snippet=snippet, - account_id=current_user.id, - ) - session.commit() + SnippetService.delete_snippet( + session=session, + snippet=snippet, + account_id=current_user.id, + ) return "", 204 @@ -314,9 +320,10 @@ class CustomizedSnippetExportApi(Resource): @edit_permission_required @rbac_permission_required(RBACCheck(RBACPermission.SNIPPETS_CREATE_AND_MODIFY, Workspace())) @with_current_tenant_id - def get(self, current_tenant_id: str, snippet_id: str): + @with_session(write=False) + def get(self, session: Session, current_tenant_id: str, snippet_id: str): """Export snippet as DSL.""" - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippet = snippet_service.get_snippet_by_id( snippet_id=snippet_id, tenant_id=current_tenant_id, @@ -328,16 +335,15 @@ class CustomizedSnippetExportApi(Resource): # Get include_secret parameter query = SnippetExportQuery.model_validate(request.args.to_dict()) - with Session(db.engine) as session: - export_service = SnippetDslService(session) - try: - result = export_service.export_snippet_dsl( - snippet=snippet, - include_secret=query.include_secret == "true", - workflow_id=query.workflow_id, - ) - except ValueError as exc: - raise NotFound(str(exc)) from exc + export_service = SnippetDslService(session) + try: + result = export_service.export_snippet_dsl( + snippet=snippet, + include_secret=query.include_secret == "true", + workflow_id=query.workflow_id, + ) + except ValueError as exc: + raise NotFound(str(exc)) from exc # Set filename with .snippet extension filename = f"{snippet.name}.snippet" @@ -432,9 +438,10 @@ class CustomizedSnippetCheckDependenciesApi(Resource): @edit_permission_required @rbac_permission_required(RBACCheck(RBACPermission.SNIPPETS_CREATE_AND_MODIFY, Workspace())) @with_current_tenant_id - def get(self, current_tenant_id: str, snippet_id: str): + @with_session(write=False) + def get(self, session: Session, current_tenant_id: str, snippet_id: str): """Check dependencies for a snippet.""" - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippet = snippet_service.get_snippet_by_id( snippet_id=snippet_id, tenant_id=current_tenant_id, @@ -443,9 +450,8 @@ class CustomizedSnippetCheckDependenciesApi(Resource): if not snippet: raise NotFound("Snippet not found") - with Session(db.engine) as session: - import_service = SnippetDslService(session) - result = import_service.check_dependencies(snippet=snippet) + import_service = SnippetDslService(session) + result = import_service.check_dependencies(snippet=snippet) return result.model_dump(mode="json"), 200 @@ -462,9 +468,10 @@ class CustomizedSnippetUseCountIncrementApi(Resource): @account_initialization_required @edit_permission_required @with_current_tenant_id - def post(self, current_tenant_id: str, snippet_id: str): + @with_session + def post(self, session: Session, current_tenant_id: str, snippet_id: str): """Increment snippet use count when it is inserted into a workflow.""" - snippet_service = _snippet_service() + snippet_service = _snippet_service(session) snippet = snippet_service.get_snippet_by_id( snippet_id=snippet_id, tenant_id=current_tenant_id, @@ -473,10 +480,7 @@ class CustomizedSnippetUseCountIncrementApi(Resource): if not snippet: raise NotFound("Snippet not found") - with Session(db.engine) as session: - snippet = session.merge(snippet) - SnippetService.increment_use_count(session=session, snippet=snippet) - session.commit() - session.refresh(snippet) + SnippetService.increment_use_count(session=session, snippet=snippet) + session.flush() return {"result": "success", "use_count": snippet.use_count}, 200 diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py index 6d762791958..a7424b2c5a5 100644 --- a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py +++ b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py @@ -3,18 +3,17 @@ from __future__ import annotations import json from datetime import datetime from inspect import unwrap -from types import SimpleNamespace from unittest.mock import Mock import pytest from flask import Flask -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from werkzeug.exceptions import HTTPException, NotFound from controllers.console.snippets import snippet_workflow as snippet_workflow_module from models.account import Account, TenantAccountRole from models.snippet import CustomizedSnippet +from models.workflow import Workflow from tests.unit_tests.model_factories import make_account @@ -40,46 +39,27 @@ def _snippet(**overrides) -> CustomizedSnippet: return CustomizedSnippet(**data) -def _workflow(**overrides) -> SimpleNamespace: - data = { - "id": "workflow-1", - "graph_dict": {"nodes": [], "edges": []}, - "features_dict": {}, - "unique_hash": "hash-1", - "version": "2024-01-01 00:00:00", - "marked_name": "v1", - "marked_comment": "first version", - "created_by_account": None, - "created_at": datetime(2024, 1, 1), - "updated_by_account": None, - "updated_at": datetime(2024, 1, 1), - "tool_published": False, - "environment_variables": [], - "conversation_variables": [], - "rag_pipeline_variables": [], - } - data.update(overrides) - workflow = SimpleNamespace(**data) - workflow.get_created_by_account = Mock(return_value=workflow.created_by_account) - workflow.get_updated_by_account = Mock(return_value=workflow.updated_by_account) - workflow.get_tool_published = Mock(return_value=workflow.tool_published) +def _workflow(**overrides) -> Workflow: + workflow = Workflow.new( + tenant_id="tenant-1", + app_id="snippet-1", + type="workflow", + version="2024-01-01 00:00:00", + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by="account-1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-1" + workflow.created_at = datetime(2024, 1, 1) + workflow.updated_at = datetime(2024, 1, 1) + for name, value in overrides.items(): + setattr(workflow, name, value) return workflow -@pytest.fixture(autouse=True) -def _patch_snippet_service_factory(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: - snippet_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - - def factory(): - try: - return snippet_workflow_module.SnippetService(snippet_workflow_module._snippet_session_maker()) - except TypeError: - return snippet_workflow_module.SnippetService() - - monkeypatch.setattr(snippet_workflow_module, "_snippet_service", factory) - monkeypatch.setattr(snippet_workflow_module, "_snippet_session_maker", lambda: snippet_session_maker) - - def test_get_snippet_requires_snippet_id(app): @snippet_workflow_module.get_snippet def view(**kwargs): @@ -127,32 +107,28 @@ def test_get_snippet_raises_not_found_when_snippet_missing(app: Flask, monkeypat view(snippet_id="snippet-1") -def test_draft_workflow_get_raises_when_missing(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_draft_workflow_get_raises_when_missing( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: snippet = _snippet() - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(get_draft_workflow=Mock(return_value=None)), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "get_draft_workflow", Mock(return_value=None)) api = snippet_workflow_module.SnippetDraftWorkflowApi() handler = unwrap(api.get) with app.test_request_context("/snippets/snippet-1/workflows/draft"): with pytest.raises(snippet_workflow_module.DraftWorkflowNotExist): - handler(api, snippet=snippet) + handler(api, unbound_session, snippet=snippet) -def test_draft_workflow_get_uses_session_aware_response_source(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow = _workflow() +def test_draft_workflow_get_uses_session_aware_response_source( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workflow = _workflow(updated_by="account-2") + sqlite_session.add_all([_account(), _account("account-2")]) + sqlite_session.commit() snippet = _snippet() - session = Mock(spec=Session) - monkeypatch.setattr(snippet_workflow_module, "db", SimpleNamespace(session=Mock(return_value=session))) - monkeypatch.setattr( - snippet_workflow_module, - "_snippet_service", - lambda: SimpleNamespace(get_draft_workflow=Mock(return_value=workflow)), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "get_draft_workflow", Mock(return_value=workflow)) monkeypatch.setattr( snippet_workflow_module.WorkflowAgentPublishService, "project_draft_bindings_to_graph", @@ -163,23 +139,19 @@ def test_draft_workflow_get_uses_session_aware_response_source(app: Flask, monke handler = unwrap(api.get) with app.test_request_context("/snippets/snippet-1/workflows/draft"): - response = handler(api, snippet=snippet) + response = handler(api, sqlite_session, snippet=snippet) assert response["id"] == "workflow-1" - workflow.get_created_by_account.assert_called_once_with(session=session) - workflow.get_updated_by_account.assert_called_once_with(session=session) - workflow.get_tool_published.assert_called_once_with(session=session) + assert response["created_by"]["id"] == "account-1" + assert response["updated_by"]["id"] == "account-2" + assert response["tool_published"] is False def test_draft_workflow_post_returns_400_for_invalid_graph(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: user = _account("account-1") snippet = _snippet() sync_draft_workflow = Mock(side_effect=ValueError("invalid graph")) - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(sync_draft_workflow=sync_draft_workflow), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "sync_draft_workflow", sync_draft_workflow) api = snippet_workflow_module.SnippetDraftWorkflowApi() handler = unwrap(api.post) @@ -207,45 +179,41 @@ def test_draft_config_returns_parallel_depth_limit(app) -> None: handler = unwrap(api.get) with app.test_request_context("/snippets/snippet-1/workflows/draft/config"): - assert handler(api, snippet=SimpleNamespace(id="snippet-1")) == {"parallel_depth_limit": 3} + assert handler(api, snippet=_snippet()) == {"parallel_depth_limit": 3} -def test_published_workflow_get_returns_none_when_not_published(app) -> None: +def test_published_workflow_get_returns_none_when_not_published(app, unbound_session: Session) -> None: api = snippet_workflow_module.SnippetPublishedWorkflowApi() handler = unwrap(api.get) with app.test_request_context("/snippets/snippet-1/workflows/publish"): - assert handler(api, snippet=SimpleNamespace(id="snippet-1", is_published=False)) is None + assert handler(api, unbound_session, snippet=_snippet(is_published=False)) is None -def test_published_workflow_get_uses_session_aware_response_source(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow = _workflow() - session = Mock(spec=Session) - snippet = SimpleNamespace(id="snippet-1", is_published=True, input_fields_list=[]) - monkeypatch.setattr(snippet_workflow_module, "db", SimpleNamespace(session=Mock(return_value=session))) - monkeypatch.setattr( - snippet_workflow_module, - "_snippet_service", - lambda: SimpleNamespace(get_published_workflow=Mock(return_value=workflow)), - ) +def test_published_workflow_get_uses_session_aware_response_source( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workflow = _workflow(updated_by="account-2") + sqlite_session.add_all([_account(), _account("account-2")]) + sqlite_session.commit() + snippet = _snippet(is_published=True) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "get_published_workflow", Mock(return_value=workflow)) api = snippet_workflow_module.SnippetPublishedWorkflowApi() handler = unwrap(api.get) with app.test_request_context("/snippets/snippet-1/workflows/publish"): - response = handler(api, snippet=snippet) + response = handler(api, sqlite_session, snippet=snippet) assert response["id"] == "workflow-1" - workflow.get_created_by_account.assert_called_once_with(session=session) - workflow.get_updated_by_account.assert_called_once_with(session=session) - workflow.get_tool_published.assert_called_once_with(session=session) + assert response["created_by"]["id"] == "account-1" + assert response["updated_by"]["id"] == "account-2" + assert response["tool_published"] is False -@pytest.mark.parametrize("sqlite_session", [(CustomizedSnippet,)], indirect=True) def test_published_workflow_post_returns_400_when_publish_fails( app: Flask, monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, sqlite_session: Session, ) -> None: user = _account("account-1") @@ -258,18 +226,13 @@ def test_published_workflow_post_returns_400_when_publish_fails( session.add(snippet) raise ValueError("No valid workflow found.") - monkeypatch.setattr(snippet_workflow_module, "db", SimpleNamespace(engine=sqlite_engine)) - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(publish_workflow=Mock(side_effect=fail_publish)), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "publish_workflow", Mock(side_effect=fail_publish)) api = snippet_workflow_module.SnippetPublishedWorkflowApi() handler = unwrap(api.post) with app.test_request_context("/snippets/snippet-1/workflows/publish", method="POST", json={}): - response, status_code = handler(api, user, snippet) + response, status_code = handler(api, sqlite_session, user, snippet) assert status_code == 400 assert response == {"message": "No valid workflow found."} @@ -277,46 +240,35 @@ def test_published_workflow_post_returns_400_when_publish_fails( assert snippet.name == "Snippet" -@pytest.mark.parametrize("sqlite_session", [(CustomizedSnippet,)], indirect=True) def test_published_workflow_post_returns_success( app: Flask, monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, sqlite_session: Session, ) -> None: user = _account("account-1") snippet = _snippet() sqlite_session.add(snippet) sqlite_session.commit() - workflow = SimpleNamespace(created_at=datetime(2026, 8, 17, 12, 0, 0)) - monkeypatch.setattr(snippet_workflow_module, "db", SimpleNamespace(engine=sqlite_engine)) - monkeypatch.setattr( - snippet_workflow_module, - "_snippet_service", - lambda: SimpleNamespace(publish_workflow=Mock(return_value=workflow)), - ) + workflow = _workflow(created_at=datetime(2026, 8, 17, 12, 0, 0)) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "publish_workflow", Mock(return_value=workflow)) api = snippet_workflow_module.SnippetPublishedWorkflowApi() handler = unwrap(api.post) with app.test_request_context("/snippets/snippet-1/workflows/publish", method="POST", json={}): - response = handler(api, user, snippet) + response = handler(api, sqlite_session, user, snippet) assert response["result"] == "success" def test_default_block_configs_delegates_to_service(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: get_default_block_configs = Mock(return_value=[{"type": "llm"}]) - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(get_default_block_configs=get_default_block_configs), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "get_default_block_configs", get_default_block_configs) api = snippet_workflow_module.SnippetDefaultBlockConfigsApi() handler = unwrap(api.get) with app.test_request_context("/snippets/snippet-1/workflows/default-workflow-block-configs"): - result = handler(api, snippet=SimpleNamespace(id="snippet-1")) + result = handler(api, snippet=_snippet()) assert result == [{"type": "llm"}] get_default_block_configs.assert_called_once() @@ -325,17 +277,14 @@ def test_default_block_configs_delegates_to_service(app: Flask, monkeypatch: pyt def test_list_published_snippet_workflows_includes_input_fields( app: Flask, monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, + sqlite_session: Session, ) -> None: - workflow = _workflow(marked_name="", marked_comment="") + workflow = _workflow() input_fields = [{"variable": "query", "type": "text"}] snippet = _snippet(input_fields=json.dumps(input_fields)) - monkeypatch.setattr(snippet_workflow_module, "db", SimpleNamespace(engine=sqlite_engine)) monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(get_all_published_workflows=Mock(return_value=([workflow], False))), + snippet_workflow_module.SnippetService, "get_all_published_workflows", Mock(return_value=([workflow], False)) ) api = snippet_workflow_module.SnippetPublishedAllWorkflowApi() @@ -345,6 +294,7 @@ def test_list_published_snippet_workflows_includes_input_fields( response = handler( api, snippet_workflow_module.SnippetWorkflowListQuery.model_validate({"page": 1, "limit": 20}), + sqlite_session, snippet=snippet, ) @@ -352,18 +302,12 @@ def test_list_published_snippet_workflows_includes_input_fields( def test_restore_published_snippet_workflow_to_draft_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow = SimpleNamespace( - unique_hash="restored-hash", - updated_at=None, - created_at=datetime(2024, 1, 1), - ) + workflow = _workflow(updated_at=None) user = _account("account-1") snippet = _snippet() monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(restore_published_workflow_to_draft=lambda **_kwargs: workflow), + snippet_workflow_module.SnippetService, "restore_published_workflow_to_draft", Mock(return_value=workflow) ) api = snippet_workflow_module.SnippetDraftWorkflowRestoreApi() @@ -376,7 +320,7 @@ def test_restore_published_snippet_workflow_to_draft_success(app: Flask, monkeyp response = handler(api, user, snippet, workflow_id="published-workflow") assert response["result"] == "success" - assert response["hash"] == "restored-hash" + assert response["hash"] == workflow.unique_hash def test_restore_published_snippet_workflow_to_draft_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: @@ -384,13 +328,9 @@ def test_restore_published_snippet_workflow_to_draft_not_found(app: Flask, monke snippet = _snippet() monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace( - restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( - snippet_workflow_module.WorkflowNotFoundError("Workflow not found") - ) - ), + snippet_workflow_module.SnippetService, + "restore_published_workflow_to_draft", + Mock(side_effect=snippet_workflow_module.WorkflowNotFoundError("Workflow not found")), ) api = snippet_workflow_module.SnippetDraftWorkflowRestoreApi() @@ -411,13 +351,9 @@ def test_restore_published_snippet_workflow_to_draft_returns_400_for_draft_sourc snippet = _snippet() monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace( - restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( - snippet_workflow_module.IsDraftWorkflowError("source workflow must be published") - ) - ), + snippet_workflow_module.SnippetService, + "restore_published_workflow_to_draft", + Mock(side_effect=snippet_workflow_module.IsDraftWorkflowError("source workflow must be published")), ) api = snippet_workflow_module.SnippetDraftWorkflowRestoreApi() @@ -441,13 +377,9 @@ def test_restore_published_snippet_workflow_to_draft_returns_400_for_invalid_gra snippet = _snippet() monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace( - restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( - ValueError("invalid snippet workflow graph") - ) - ), + snippet_workflow_module.SnippetService, + "restore_published_workflow_to_draft", + Mock(side_effect=ValueError("invalid snippet workflow graph")), ) api = snippet_workflow_module.SnippetDraftWorkflowRestoreApi() @@ -464,13 +396,12 @@ def test_restore_published_snippet_workflow_to_draft_returns_400_for_invalid_gra assert exc.value.description == "invalid snippet workflow graph" -@pytest.mark.parametrize("sqlite_session", [(CustomizedSnippet,)], indirect=True) def test_update_published_snippet_workflow_returns_updated_workflow( app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, ) -> None: - workflow = _workflow() + workflow = _workflow(marked_name="v1", marked_comment="first version") user = _account("account-1") input_fields = [{"variable": "query", "type": "text"}] snippet = _snippet(input_fields=json.dumps(input_fields)) @@ -483,11 +414,7 @@ def test_update_published_snippet_workflow_returns_updated_workflow( return workflow update_workflow = Mock(side_effect=update_persisted_snippet) - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(update_workflow=update_workflow), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "update_workflow", update_workflow) api = snippet_workflow_module.SnippetWorkflowByIdApi() handler = unwrap(api.patch) @@ -502,6 +429,7 @@ def test_update_published_snippet_workflow_returns_updated_workflow( snippet_workflow_module.WorkflowUpdatePayload.model_validate( {"marked_name": "v1", "marked_comment": "first version"} ), + sqlite_session, user, snippet, workflow_id="workflow-1", @@ -521,7 +449,7 @@ def test_update_published_snippet_workflow_returns_updated_workflow( assert snippet.description == "Updated in transaction" -def test_update_published_snippet_workflow_returns_400_when_no_fields(app: Flask) -> None: +def test_update_published_snippet_workflow_returns_400_when_no_fields(app: Flask, unbound_session: Session) -> None: api = snippet_workflow_module.SnippetWorkflowByIdApi() handler = unwrap(api.patch) @@ -529,6 +457,7 @@ def test_update_published_snippet_workflow_returns_400_when_no_fields(app: Flask response, status_code = handler( api, snippet_workflow_module.WorkflowUpdatePayload(), + unbound_session, _account("account-1"), _snippet(), workflow_id="workflow-1", @@ -538,7 +467,6 @@ def test_update_published_snippet_workflow_returns_400_when_no_fields(app: Flask assert response == {"message": "No valid fields to update"} -@pytest.mark.parametrize("sqlite_session", [(CustomizedSnippet,)], indirect=True) def test_update_published_snippet_workflow_raises_not_found( app: Flask, monkeypatch: pytest.MonkeyPatch, @@ -554,9 +482,7 @@ def test_update_published_snippet_workflow_raises_not_found( merged_snippet.name = "Rolled back name" monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(update_workflow=Mock(side_effect=update_missing_workflow)), + snippet_workflow_module.SnippetService, "update_workflow", Mock(side_effect=update_missing_workflow) ) api = snippet_workflow_module.SnippetWorkflowByIdApi() @@ -571,11 +497,13 @@ def test_update_published_snippet_workflow_raises_not_found( handler( api, snippet_workflow_module.WorkflowUpdatePayload.model_validate({"marked_name": "v1"}), + sqlite_session, user, snippet, workflow_id="missing-workflow", ) + sqlite_session.rollback() sqlite_session.refresh(snippet) assert snippet.name == "Snippet" @@ -583,11 +511,7 @@ def test_update_published_snippet_workflow_raises_not_found( def test_delete_published_snippet_workflow_succeeds(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: snippet = _snippet() delete_workflow = Mock(return_value=True) - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(delete_workflow=delete_workflow), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "delete_workflow", delete_workflow) api = snippet_workflow_module.SnippetWorkflowByIdApi() handler = unwrap(api.delete) @@ -609,9 +533,7 @@ def test_delete_published_snippet_workflow_raises_not_found(app: Flask, monkeypa raise ValueError("Workflow with ID missing-workflow not found") monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(delete_workflow=Mock(side_effect=delete_missing_workflow)), + snippet_workflow_module.SnippetService, "delete_workflow", Mock(side_effect=delete_missing_workflow) ) api = snippet_workflow_module.SnippetWorkflowByIdApi() @@ -629,9 +551,7 @@ def test_delete_published_snippet_workflow_raises_bad_request_when_in_use( raise snippet_workflow_module.WorkflowInUseError("Cannot delete workflow that is currently in use") monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(delete_workflow=Mock(side_effect=delete_active_workflow)), + snippet_workflow_module.SnippetService, "delete_workflow", Mock(side_effect=delete_active_workflow) ) api = snippet_workflow_module.SnippetWorkflowByIdApi() @@ -648,11 +568,7 @@ def test_workflow_run_detail_raises_not_found_when_run_missing( app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: snippet = _snippet() - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace(get_snippet_workflow_run=Mock(return_value=None)), - ) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "get_snippet_workflow_run", Mock(return_value=None)) api = snippet_workflow_module.SnippetWorkflowRunDetailApi() handler = unwrap(api.get) @@ -666,15 +582,9 @@ def test_draft_node_last_run_raises_not_found_when_execution_missing( app: Flask, monkeypatch: pytest.MonkeyPatch ) -> None: snippet = _snippet() - draft_workflow = SimpleNamespace(id="workflow-1") - monkeypatch.setattr( - snippet_workflow_module, - "SnippetService", - lambda: SimpleNamespace( - get_draft_workflow=Mock(return_value=draft_workflow), - get_snippet_node_last_run=Mock(return_value=None), - ), - ) + draft_workflow = _workflow(version=Workflow.VERSION_DRAFT) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "get_draft_workflow", Mock(return_value=draft_workflow)) + monkeypatch.setattr(snippet_workflow_module.SnippetService, "get_snippet_node_last_run", Mock(return_value=None)) api = snippet_workflow_module.SnippetDraftNodeLastRunApi() handler = unwrap(api.get) @@ -695,14 +605,14 @@ def test_workflow_task_stop_uses_queue_flag_and_graph_command(app: Flask, monkey monkeypatch.setattr( snippet_workflow_module, "GraphEngineManager", - Mock(return_value=SimpleNamespace(send_stop_command=send_stop_command)), + Mock(return_value=Mock(send_stop_command=send_stop_command)), ) api = snippet_workflow_module.SnippetWorkflowTaskStopApi() handler = unwrap(api.post) with app.test_request_context("/snippets/snippet-1/workflow-runs/tasks/task-1/stop", method="POST"): - result = handler(api, snippet=SimpleNamespace(id="snippet-1"), task_id="task-1") + result = handler(api, snippet=_snippet(), task_id="task-1") assert result == {"result": "success"} set_stop_flag.assert_called_once_with("task-1") diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py index fb41206f65c..578b2b1c03b 100644 --- a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py +++ b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py @@ -1,28 +1,17 @@ -from collections.abc import Iterator +import json from inspect import unwrap -from types import SimpleNamespace from unittest.mock import Mock import pytest from flask import Flask -from sqlalchemy import event, select -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy import select +from sqlalchemy.orm import Session from controllers.console.snippets import snippet_workflow_draft_variable as module -from graphon.variables import StringSegment -from graphon.variables.types import SegmentType +from graphon.variables import StringSegment, StringVariable from models.account import Account, AccountStatus -from models.workflow import WorkflowDraftVariable, WorkflowDraftVariableFile - -pytestmark = [ - pytest.mark.usefixtures("sqlite_session"), - pytest.mark.parametrize( - "sqlite_session", - [(WorkflowDraftVariable, WorkflowDraftVariableFile)], - indirect=True, - ), -] +from models.snippet import CustomizedSnippet +from models.workflow import Workflow, WorkflowDraftVariable def _make_account() -> Account: @@ -35,6 +24,34 @@ def _make_account() -> Account: return account +def _make_snippet(snippet_id: str = "snippet-1") -> CustomizedSnippet: + return CustomizedSnippet( + id=snippet_id, + tenant_id="tenant-1", + name="Snippet", + description="Description", + type="node", + created_by="user-1", + ) + + +def _make_workflow(*, environment_variables: list[StringVariable] | None = None) -> Workflow: + workflow = Workflow.new( + tenant_id="tenant-1", + app_id="snippet-1", + type="workflow", + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by="user-1", + environment_variables=environment_variables or [], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-1" + return workflow + + def _make_node_variable( variable_id: str, *, @@ -58,17 +75,6 @@ def _make_node_variable( return variable -@pytest.fixture(autouse=True) -def _patch_snippet_service_factory(monkeypatch: pytest.MonkeyPatch) -> None: - def factory(): - service_factory = module.SnippetService - if isinstance(service_factory, type): - return service_factory.__new__(service_factory) - return service_factory() - - monkeypatch.setattr(module, "_snippet_service", factory) - - @pytest.fixture def app() -> Flask: app = Flask("test_snippet_workflow_draft_variable") @@ -76,28 +82,13 @@ def app() -> Flask: return app -@pytest.fixture -def controller_sessions( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, -) -> Iterator[scoped_session[Session]]: - """Bind both controller session styles to the isolated SQLite engine.""" - sessions = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine, session=sessions)) - try: - yield sessions - finally: - sessions.remove() - - def _persist_variables(sqlite_session: Session, *variables: WorkflowDraftVariable) -> None: sqlite_session.add_all(variables) sqlite_session.commit() -def _variable_ids(sqlite_engine: Engine) -> set[str]: - with Session(sqlite_engine) as session: - return set(session.scalars(select(WorkflowDraftVariable.id))) +def _variable_ids(session: Session) -> set[str]: + return set(session.scalars(select(WorkflowDraftVariable.id))) def test_ensure_snippet_draft_variable_row_allowed_rejects_system_variable() -> None: @@ -137,7 +128,7 @@ def test_conversation_variables_returns_empty_list(app: Flask) -> None: handler = unwrap(api.get) with app.test_request_context("/"): - result = handler(api, _make_account(), snippet=SimpleNamespace(id="snippet-1")) + result = handler(api, _make_account(), snippet=_make_snippet()) assert result == {"items": []} @@ -147,7 +138,7 @@ def test_system_variables_returns_empty_list(app: Flask) -> None: handler = unwrap(api.get) with app.test_request_context("/"): - result = handler(api, _make_account(), snippet=SimpleNamespace(id="snippet-1")) + result = handler(api, _make_account(), snippet=_make_snippet()) assert result == {"items": []} @@ -155,8 +146,6 @@ def test_system_variables_returns_empty_list(app: Flask) -> None: def test_delete_variable_collection_deletes_only_current_user_variables( app: Flask, sqlite_session: Session, - sqlite_engine: Engine, - controller_sessions: scoped_session[Session], ) -> None: matching = _make_node_variable("matching", name="matching") matching_second = _make_node_variable("matching-second", node_id="tool-1", name="matching-second") @@ -167,22 +156,18 @@ def test_delete_variable_collection_deletes_only_current_user_variables( handler = unwrap(api.delete) with app.test_request_context("/", method="DELETE"): - response = handler(api, _make_account(), snippet=SimpleNamespace(id="snippet-1")) + response = handler(api, sqlite_session, _make_account(), snippet=_make_snippet()) assert response.status_code == 204 - assert _variable_ids(sqlite_engine) == {other_user.id, other_snippet.id} - assert not controller_sessions().in_transaction() + assert _variable_ids(sqlite_session) == {other_user.id, other_snippet.id} def test_variable_collection_get_raises_when_draft_workflow_missing( app: Flask, monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, ) -> None: - monkeypatch.setattr( - module, - "SnippetService", - Mock(return_value=SimpleNamespace(get_draft_workflow=Mock(return_value=None))), - ) + monkeypatch.setattr(module.SnippetService, "get_draft_workflow", Mock(return_value=None)) api = module.SnippetWorkflowVariableCollectionApi() handler = unwrap(api.get) @@ -191,16 +176,16 @@ def test_variable_collection_get_raises_when_draft_workflow_missing( with pytest.raises(module.DraftWorkflowNotExist): handler( api, + sqlite_session, module.WorkflowDraftVariableListQuery(page=1, limit=20), _make_account(), - snippet=SimpleNamespace(id="snippet-1"), + snippet=_make_snippet(), ) def test_node_variable_collection_get_lists_persisted_node_variables( app: Flask, sqlite_session: Session, - controller_sessions: scoped_session[Session], ) -> None: matching = _make_node_variable("matching", name="matching") other_node = _make_node_variable("other-node", node_id="tool-1", name="other-node") @@ -211,17 +196,14 @@ def test_node_variable_collection_get_lists_persisted_node_variables( handler = unwrap(api.get) with app.test_request_context("/"): - result = handler(api, _make_account(), snippet=SimpleNamespace(id="snippet-1"), node_id="llm-1") + result = handler(api, sqlite_session, _make_account(), snippet=_make_snippet(), node_id="llm-1") assert [variable["id"] for variable in result["items"]] == [matching.id] - assert controller_sessions().get_bind() is not None def test_node_variable_collection_delete_deletes_only_requested_node_variables( app: Flask, sqlite_session: Session, - sqlite_engine: Engine, - controller_sessions: scoped_session[Session], ) -> None: matching = _make_node_variable("matching", name="matching") matching_second = _make_node_variable("matching-second", name="matching-second") @@ -232,51 +214,36 @@ def test_node_variable_collection_delete_deletes_only_requested_node_variables( handler = unwrap(api.delete) with app.test_request_context("/", method="DELETE"): - response = handler(api, _make_account(), snippet=SimpleNamespace(id="snippet-1"), node_id="llm-1") + response = handler(api, sqlite_session, _make_account(), snippet=_make_snippet(), node_id="llm-1") assert response.status_code == 204 - assert _variable_ids(sqlite_engine) == {other_node.id, other_user.id} - assert not controller_sessions().in_transaction() + assert _variable_ids(sqlite_session) == {other_node.id, other_user.id} def test_variable_patch_returns_persisted_variable_without_committing_when_no_changes( app: Flask, sqlite_session: Session, - controller_sessions: scoped_session[Session], ) -> None: variable = _make_node_variable("var-1") _persist_variables(sqlite_session, variable) - session = controller_sessions() - commits: list[bool] = [] - - def record_commit(_session: Session) -> None: - commits.append(True) - - event.listen(session, "after_commit", record_commit) api = module.SnippetVariableApi() handler = unwrap(api.patch) - try: - with app.test_request_context("/", method="PATCH", json={}): - result = handler( - api, - module.WorkflowDraftVariableUpdatePayload(), - _make_account(), - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), - variable_id="var-1", - ) - finally: - event.remove(session, "after_commit", record_commit) + with app.test_request_context("/", method="PATCH", json={}): + result = handler( + api, + module.WorkflowDraftVariableUpdatePayload(), + sqlite_session, + _make_account(), + snippet=_make_snippet(), + variable_id="var-1", + ) assert result["id"] == variable.id - assert commits == [] - assert session.in_transaction() def test_variable_delete_deletes_persisted_variable( app: Flask, sqlite_session: Session, - sqlite_engine: Engine, - controller_sessions: scoped_session[Session], ) -> None: variable = _make_node_variable("var-1") retained = _make_node_variable("var-2", name="retained") @@ -287,73 +254,60 @@ def test_variable_delete_deletes_persisted_variable( with app.test_request_context("/", method="DELETE"): response = handler( api, + sqlite_session, _make_account(), - snippet=SimpleNamespace(id="snippet-1"), + snippet=_make_snippet(), variable_id=variable.id, ) assert response.status_code == 204 - assert _variable_ids(sqlite_engine) == {retained.id} - assert not controller_sessions().in_transaction() + assert _variable_ids(sqlite_session) == {retained.id} def test_variable_reset_deletes_variable_without_node_execution( app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, - sqlite_engine: Engine, - controller_sessions: scoped_session[Session], ) -> None: variable = _make_node_variable("var-1", node_execution_id=None) _persist_variables(sqlite_session, variable) - monkeypatch.setattr( - module, - "SnippetService", - Mock(return_value=SimpleNamespace(get_draft_workflow=Mock(return_value=SimpleNamespace(id="workflow-1")))), - ) + monkeypatch.setattr(module.SnippetService, "get_draft_workflow", Mock(return_value=_make_workflow())) api = module.SnippetVariableResetApi() handler = unwrap(api.put) with app.test_request_context("/", method="PUT"): response = handler( api, + sqlite_session, _make_account(), - snippet=SimpleNamespace(id="snippet-1"), + snippet=_make_snippet(), variable_id=variable.id, ) assert response.status_code == 204 - assert _variable_ids(sqlite_engine) == set() - assert not controller_sessions().in_transaction() + assert _variable_ids(sqlite_session) == set() def test_environment_variables_returns_workflow_environment_variables( app: Flask, monkeypatch: pytest.MonkeyPatch, ) -> None: - env_var = SimpleNamespace( + env_var = StringVariable( id="env-1", name="API_KEY", description="secret", selector=["env", "API_KEY"], - value_type=SegmentType.SECRET, value="sk-test", ) monkeypatch.setattr( - module, - "SnippetService", - Mock( - return_value=SimpleNamespace( - get_draft_workflow=Mock(return_value=SimpleNamespace(environment_variables=[env_var])) - ) - ), + module.SnippetService, "get_draft_workflow", Mock(return_value=_make_workflow(environment_variables=[env_var])) ) api = module.SnippetEnvironmentVariableCollectionApi() handler = unwrap(api.get) with app.test_request_context("/"): - result = handler(api, _make_account(), snippet=SimpleNamespace(id="snippet-1")) + result = handler(api, _make_account(), snippet=_make_snippet()) assert result == { "items": [ @@ -363,7 +317,7 @@ def test_environment_variables_returns_workflow_environment_variables( "name": "API_KEY", "description": "secret", "selector": ["env", "API_KEY"], - "value_type": "secret", + "value_type": "string", "value": "sk-test", "edited": False, "visible": True, diff --git a/api/tests/unit_tests/controllers/console/workspace/test_snippets.py b/api/tests/unit_tests/controllers/console/workspace/test_snippets.py index 78943e87482..1ceba157337 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_snippets.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_snippets.py @@ -6,8 +6,7 @@ from unittest.mock import Mock import pytest from flask import Flask from pydantic import ValidationError -from sqlalchemy import Engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session from werkzeug.exceptions import BadRequest, NotFound from controllers.console.workspace import snippets as snippets_module @@ -17,22 +16,6 @@ from services.snippet_dsl_service import ImportStatus, SnippetImportInfo from tests.unit_tests.model_factories import make_account -@pytest.fixture(autouse=True) -def _patch_snippet_service_factory( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], -): - def factory(): - return snippets_module.SnippetService.__new__(snippets_module.SnippetService) - - database_session = scoped_session(sqlite_session_factory) - monkeypatch.setattr(snippets_module, "_snippet_service", factory) - monkeypatch.setattr(snippets_module, "db", SimpleNamespace(engine=sqlite_engine, session=database_session)) - yield - database_session.remove() - - def _account(account_id: str = "account-1") -> Account: return make_account( account_id=account_id, @@ -437,7 +420,7 @@ def _persisted_name(session: Session) -> str | None: return stored.name if stored else None -def test_delete_snippet_deletes_and_commits(app: Flask, monkeypatch: pytest.MonkeyPatch): +def test_delete_snippet_delegates_to_service(app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): snippet = _snippet() user = _account() delete_snippet = Mock() @@ -449,16 +432,16 @@ def test_delete_snippet_deletes_and_commits(app: Flask, monkeypatch: pytest.Monk handler = unwrap(api.delete) with app.test_request_context("/workspaces/current/customized-snippets/snippet-1", method="DELETE"): - response, status_code = handler(api, "tenant-1", user, snippet_id="snippet-1") + response, status_code = handler(api, sqlite_session, "tenant-1", user, snippet_id="snippet-1") assert status_code == 204 assert response == "" assert delete_snippet.call_args.kwargs["account_id"] == user.id - assert isinstance(delete_snippet.call_args.kwargs["session"], Session) + assert delete_snippet.call_args.kwargs["session"] is sqlite_session assert isinstance(delete_snippet.call_args.kwargs["snippet"], CustomizedSnippet) -def test_export_snippet_returns_yaml_attachment(app: Flask, monkeypatch: pytest.MonkeyPatch): +def test_export_snippet_returns_yaml_attachment(app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): snippet = _snippet(name="Snippet One") export_snippet_dsl = Mock(return_value="version: 0.1.0\nkind: snippet\n") @@ -475,7 +458,7 @@ def test_export_snippet_returns_yaml_attachment(app: Flask, monkeypatch: pytest. with app.test_request_context( "/workspaces/current/customized-snippets/snippet-1/export?include_secret=true&workflow_id=workflow-1" ): - response = handler(api, "tenant-1", snippet_id="snippet-1") + response = handler(api, sqlite_session, "tenant-1", snippet_id="snippet-1") assert response.status_code == 200 assert response.get_data(as_text=True) == "version: 0.1.0\nkind: snippet\n" @@ -484,7 +467,9 @@ def test_export_snippet_returns_yaml_attachment(app: Flask, monkeypatch: pytest. export_snippet_dsl.assert_called_once_with(snippet=snippet, include_secret=True, workflow_id="workflow-1") -def test_export_snippet_raises_not_found_for_missing_workflow(app: Flask, monkeypatch: pytest.MonkeyPatch): +def test_export_snippet_raises_not_found_for_missing_workflow( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +): snippet = _snippet(name="Snippet One") monkeypatch.setattr(snippets_module.SnippetService, "get_snippet_by_id", Mock(return_value=snippet)) @@ -502,7 +487,7 @@ def test_export_snippet_raises_not_found_for_missing_workflow(app: Flask, monkey with app.test_request_context("/workspaces/current/customized-snippets/snippet-1/export?workflow_id=workflow-1"): with pytest.raises(NotFound, match="Missing published workflow workflow-1"): - handler(api, "tenant-1", snippet_id="snippet-1") + handler(api, sqlite_session, "tenant-1", snippet_id="snippet-1") def test_import_snippet_returns_202_for_pending_confirmation( @@ -588,7 +573,9 @@ def test_import_confirm_returns_200_for_completed_import( confirm_import.assert_called_once_with(import_id="import-1", account=user) -def test_check_dependencies_raises_when_snippet_missing(app: Flask, monkeypatch: pytest.MonkeyPatch): +def test_check_dependencies_raises_when_snippet_missing( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +): monkeypatch.setattr(snippets_module.SnippetService, "get_snippet_by_id", Mock(return_value=None)) api = snippets_module.CustomizedSnippetCheckDependenciesApi() @@ -596,10 +583,12 @@ def test_check_dependencies_raises_when_snippet_missing(app: Flask, monkeypatch: with app.test_request_context("/workspaces/current/customized-snippets/snippet-1/check-dependencies"): with pytest.raises(NotFound, match="Snippet not found"): - handler(api, "tenant-1", snippet_id="snippet-1") + handler(api, sqlite_session, "tenant-1", snippet_id="snippet-1") -def test_check_dependencies_returns_dependency_result(app: Flask, monkeypatch: pytest.MonkeyPatch): +def test_check_dependencies_returns_dependency_result( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +): snippet = _snippet() check_dependencies = Mock(return_value=SimpleNamespace(model_dump=Mock(return_value={"leaked_dependencies": []}))) @@ -614,14 +603,16 @@ def test_check_dependencies_returns_dependency_result(app: Flask, monkeypatch: p handler = unwrap(api.get) with app.test_request_context("/workspaces/current/customized-snippets/snippet-1/check-dependencies"): - response, status_code = handler(api, "tenant-1", snippet_id="snippet-1") + response, status_code = handler(api, sqlite_session, "tenant-1", snippet_id="snippet-1") assert status_code == 200 assert response == {"leaked_dependencies": []} check_dependencies.assert_called_once_with(snippet=snippet) -def test_increment_use_count_raises_when_snippet_missing(app: Flask, monkeypatch: pytest.MonkeyPatch): +def test_increment_use_count_raises_when_snippet_missing( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +): monkeypatch.setattr(snippets_module.SnippetService, "get_snippet_by_id", Mock(return_value=None)) api = snippets_module.CustomizedSnippetUseCountIncrementApi() @@ -632,10 +623,12 @@ def test_increment_use_count_raises_when_snippet_missing(app: Flask, monkeypatch method="POST", ): with pytest.raises(NotFound, match="Snippet not found"): - handler(api, "tenant-1", snippet_id="snippet-1") + handler(api, sqlite_session, "tenant-1", snippet_id="snippet-1") -def test_increment_use_count_returns_refreshed_count(app: Flask, monkeypatch: pytest.MonkeyPatch): +def test_increment_use_count_returns_refreshed_count( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +): snippet = _snippet(use_count=2) def increment_use_count(*, session: Session, snippet: CustomizedSnippet) -> None: @@ -653,9 +646,9 @@ def test_increment_use_count_returns_refreshed_count(app: Flask, monkeypatch: py "/workspaces/current/customized-snippets/snippet-1/use-count/increment", method="POST", ): - response, status_code = handler(api, "tenant-1", snippet_id="snippet-1") + response, status_code = handler(api, sqlite_session, "tenant-1", snippet_id="snippet-1") assert status_code == 200 assert response == {"result": "success", "use_count": 3} - assert isinstance(increment_use_count_mock.call_args.kwargs["session"], Session) + assert increment_use_count_mock.call_args.kwargs["session"] is sqlite_session assert isinstance(increment_use_count_mock.call_args.kwargs["snippet"], CustomizedSnippet)