mtfyun/dify
1
0 Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-11-13 14:51:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: migrate console snippets sessions and ORM models to SQLite (#40565)

Browse Source
This commit is contained in:
Asuka Minato authored and GitHub committed 2026-09-19 09:26:43 +00:00
1 parent ceacf275b1
commit 3a9636faad
6 files changed
+361 -482
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

api/controllers/console/snippets/snippet_workflow.py
+74 -67
View File
@@ -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/<uuid:snippet_id>/workflows/draft/nodes/<string:node_id>/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/<uuid:snippet_id>/workflows/draft/iteration/nodes/<string:node_id>/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)
api/controllers/console/snippets/snippet_workflow_draft_variable.py
+54 -43
View File
@@ -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)
api/controllers/console/workspace/snippets.py
+47 -43
View File
@@ -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
api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py
+92 -182
View File
@@ -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")
api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py
+67 -113
View File
@@ -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,
api/tests/unit_tests/controllers/console/workspace/test_snippets.py
+27 -34
View File
@@ -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)