From f3e5f41a3c57506aa99a1227288e5f1f909d9deb Mon Sep 17 00:00:00 2001 From: Keith <148296621+keith991001@users.noreply.github.com> Date: Tue, 8 Sep 2026 09:31:08 +0000 Subject: [PATCH] refactor(models): pass session into WorkflowNodeExecutionModel accessors (#41968) --- api/controllers/console/app/workflow.py | 11 +- .../rag_pipeline/rag_pipeline_workflow.py | 9 +- .../console/snippets/snippet_workflow.py | 7 +- api/fields/workflow_run_fields.py | 25 ++++ api/models/workflow.py | 10 +- api/services/rag_pipeline/rag_pipeline.py | 4 +- api/services/snippet_service.py | 5 +- .../workflow_node_execution_trace_service.py | 23 ++-- api/services/workflow_run_service.py | 3 +- .../test_workflow_node_execution_model.py | 8 +- .../console/app/test_workflow_run_api.py | 5 +- .../test_rag_pipeline_workflow_apis.py | 7 +- .../test_workflow_node_execution_accessors.py | 110 ++++++++++++++++++ .../rag_pipeline/test_rag_pipeline_service.py | 3 +- .../services/test_snippet_service.py | 6 +- ...t_workflow_node_execution_trace_service.py | 30 +++-- .../services/test_workflow_run_service.py | 4 +- 17 files changed, 222 insertions(+), 48 deletions(-) create mode 100644 api/tests/unit_tests/models/test_workflow_node_execution_accessors.py diff --git a/api/controllers/console/app/workflow.py b/api/controllers/console/app/workflow.py index aa596e39a8e..23cc11b8282 100644 --- a/api/controllers/console/app/workflow.py +++ b/api/controllers/console/app/workflow.py @@ -77,7 +77,10 @@ from factories import file_factory, variable_factory from fields.base import ResponseModel from fields.conversation_variable_fields import WorkflowConversationVariableResponse from fields.member_fields import SimpleAccount -from fields.workflow_run_fields import WorkflowRunNodeExecutionResponse +from fields.workflow_run_fields import ( + WorkflowRunNodeExecutionResponse, + node_execution_response_source, +) from graphon.enums import NodeType from graphon.file import File from graphon.file import helpers as file_helpers @@ -1260,7 +1263,7 @@ class DraftWorkflowNodeRunApi(Resource): ) return WorkflowRunNodeExecutionResponse.model_validate( - workflow_node_execution, from_attributes=True + node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True ).model_dump(mode="json") @@ -1680,7 +1683,9 @@ class DraftWorkflowNodeLastRunApi(Resource): ) if node_exec is None: raise NotFound("last run not found") - return WorkflowRunNodeExecutionResponse.model_validate(node_exec, from_attributes=True).model_dump(mode="json") + return WorkflowRunNodeExecutionResponse.model_validate( + node_execution_response_source(node_exec, session=db.session()), from_attributes=True + ).model_dump(mode="json") @console_ns.route("/apps//workflows/draft/trigger/run") diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py index a8e98298eb5..230cc5a3cd5 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py @@ -53,6 +53,7 @@ from fields.workflow_run_fields import ( WorkflowRunNodeExecutionListResponse, WorkflowRunNodeExecutionResponse, WorkflowRunPaginationResponse, + node_execution_response_source, ) from graphon.model_runtime.utils.encoders import jsonable_encoder from libs import helper @@ -506,7 +507,7 @@ class RagPipelineDraftNodeRunApi(Resource): raise ValueError("Workflow node execution not found") return WorkflowRunNodeExecutionResponse.model_validate( - workflow_node_execution, from_attributes=True + node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True ).model_dump(mode="json") @@ -1023,7 +1024,9 @@ class RagPipelineWorkflowLastRunApi(Resource): ) if node_exec is None: raise NotFound("last run not found") - return WorkflowRunNodeExecutionResponse.model_validate(node_exec, from_attributes=True).model_dump(mode="json") + return WorkflowRunNodeExecutionResponse.model_validate( + node_execution_response_source(node_exec, session=db.session()), from_attributes=True + ).model_dump(mode="json") @console_ns.route("/rag/pipelines/transform/datasets/") @@ -1085,7 +1088,7 @@ class RagPipelineDatasourceVariableApi(Resource): current_user=current_user, ) return WorkflowRunNodeExecutionResponse.model_validate( - workflow_node_execution, from_attributes=True + node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True ).model_dump(mode="json") diff --git a/api/controllers/console/snippets/snippet_workflow.py b/api/controllers/console/snippets/snippet_workflow.py index c90d75384f8..5cc8d42d5a6 100644 --- a/api/controllers/console/snippets/snippet_workflow.py +++ b/api/controllers/console/snippets/snippet_workflow.py @@ -51,6 +51,7 @@ from fields.workflow_run_fields import ( WorkflowRunNodeExecutionListResponse, WorkflowRunNodeExecutionResponse, WorkflowRunPaginationResponse, + node_execution_response_source, ) from graphon.graph_engine.manager import GraphEngineManager from libs import helper @@ -635,7 +636,7 @@ class SnippetDraftNodeRunApi(Resource): ) return WorkflowRunNodeExecutionResponse.model_validate( - workflow_node_execution, from_attributes=True + node_execution_response_source(workflow_node_execution, session=db.session()), from_attributes=True ).model_dump(mode="json") @@ -672,7 +673,9 @@ class SnippetDraftNodeLastRunApi(Resource): if node_exec is None: raise NotFound("Node last run not found") - return WorkflowRunNodeExecutionResponse.model_validate(node_exec, from_attributes=True).model_dump(mode="json") + return WorkflowRunNodeExecutionResponse.model_validate( + node_execution_response_source(node_exec, session=db.session()), from_attributes=True + ).model_dump(mode="json") @console_ns.route("/snippets//workflows/draft/iteration/nodes//run") diff --git a/api/fields/workflow_run_fields.py b/api/fields/workflow_run_fields.py index 56b0cd248b6..34c43138b73 100644 --- a/api/fields/workflow_run_fields.py +++ b/api/fields/workflow_run_fields.py @@ -7,11 +7,13 @@ kept only for workflow app-log endpoints that still build legacy log models. from __future__ import annotations +from dataclasses import dataclass from datetime import datetime from typing import Any from flask_restx import Namespace, fields from pydantic import AliasChoices, Field, field_validator +from sqlalchemy.orm import Session from fields.base import ResponseModel from fields.end_user_fields import SimpleEndUser @@ -187,5 +189,28 @@ class WorkflowRunNodeExecutionResponse(ResponseModel): return to_timestamp(value) +@dataclass(frozen=True) +class WorkflowNodeExecutionResponseSource: + """Expose session-backed node-execution accessors during response validation.""" + + node_execution: Any + session: Session + + @property + def created_by_account(self) -> Any: + return self.node_execution.created_by_account(self.session) + + @property + def created_by_end_user(self) -> Any: + return self.node_execution.created_by_end_user(self.session) + + def __getattr__(self, name: str) -> Any: + return getattr(self.node_execution, name) # guard-ignore: no-new-getattr -- delegates model fields + + +def node_execution_response_source(node_execution: Any, *, session: Session) -> WorkflowNodeExecutionResponseSource: + return WorkflowNodeExecutionResponseSource(node_execution=node_execution, session=session) + + class WorkflowRunNodeExecutionListResponse(ResponseModel): data: list[WorkflowRunNodeExecutionResponse] diff --git a/api/models/workflow.py b/api/models/workflow.py index 95159da2bd8..dbc248e5a11 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -1100,22 +1100,20 @@ class WorkflowNodeExecutionModel(Base): # This model is expected to have `offlo ) ) - @property - def created_by_account(self): + def created_by_account(self, session: orm.Session) -> Account | None: created_by_role = CreatorUserRole(self.created_by_role) if created_by_role == CreatorUserRole.ACCOUNT: stmt = select(Account).where(Account.id == self.created_by) - return db.session.scalar(stmt) + return session.scalar(stmt) return None - @property - def created_by_end_user(self): + def created_by_end_user(self, session: orm.Session): from .model import EndUser created_by_role = CreatorUserRole(self.created_by_role) if created_by_role == CreatorUserRole.END_USER: stmt = select(EndUser).where(EndUser.id == self.created_by) - return db.session.scalar(stmt) + return session.scalar(stmt) return None @property diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index 00daaa2bdd0..c5f1e1935ca 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -1266,7 +1266,9 @@ class RagPipelineService: app_id=pipeline.id, workflow_run_id=run_id, ) - return assemble_workflow_node_execution_traces(node_executions, self._node_execution_service_repo) + return assemble_workflow_node_execution_traces( + node_executions, self._node_execution_service_repo, session=self._session + ) @staticmethod def publish_customized_pipeline_template( diff --git a/api/services/snippet_service.py b/api/services/snippet_service.py index cf0809dd520..51ca4bf9369 100644 --- a/api/services/snippet_service.py +++ b/api/services/snippet_service.py @@ -989,7 +989,10 @@ class SnippetService: workflow_run_id=workflow_run.id, ) - return assemble_workflow_node_execution_traces(node_executions, self._node_execution_service_repo) + with self._session_scope() as session: + return assemble_workflow_node_execution_traces( + node_executions, self._node_execution_service_repo, session=session + ) # --- Node Execution Operations --- diff --git a/api/services/workflow_node_execution_trace_service.py b/api/services/workflow_node_execution_trace_service.py index 4d028318c81..f016afa80ea 100644 --- a/api/services/workflow_node_execution_trace_service.py +++ b/api/services/workflow_node_execution_trace_service.py @@ -8,6 +8,7 @@ from enum import Enum from typing import Any from pydantic import BaseModel, ConfigDict, ValidationError +from sqlalchemy.orm import Session from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt from libs.helper import to_timestamp @@ -58,23 +59,27 @@ class WorkflowNodeExecutionTrace(BaseModel): def assemble_workflow_node_execution_traces( executions: Sequence[WorkflowNodeExecutionModel], repository: DifyAPIWorkflowNodeExecutionRepository, + *, + session: Session, ) -> list[WorkflowNodeExecutionTrace]: """Expand valid persisted retry attempts before each terminal execution.""" traces: list[WorkflowNodeExecutionTrace] = [] for execution in executions: - traces.extend(_expand_execution(execution, repository)) + traces.extend(_expand_execution(execution, repository, session=session)) return traces def _expand_execution( execution: WorkflowNodeExecutionModel, repository: DifyAPIWorkflowNodeExecutionRepository, + *, + session: Session, ) -> list[WorkflowNodeExecutionTrace]: full_process_data = _load_full_process_data(execution, repository) retry_attempts = _parse_retry_attempts(execution, full_process_data) terminal_metadata = execution.execution_metadata_dict - traces = [_retry_trace(execution, attempt, terminal_metadata) for attempt in retry_attempts] - traces.append(_terminal_trace(execution)) + traces = [_retry_trace(execution, attempt, terminal_metadata, session=session) for attempt in retry_attempts] + traces.append(_terminal_trace(execution, session=session)) return traces @@ -133,6 +138,8 @@ def _retry_trace( execution: WorkflowNodeExecutionModel, attempt: WorkflowNodeRetryAttempt, terminal_metadata: Mapping[str, Any], + *, + session: Session, ) -> WorkflowNodeExecutionTrace: truncator = VariableTruncator.default() inputs, inputs_truncated = truncator.truncate_variable_mapping(attempt.inputs) @@ -163,8 +170,8 @@ def _retry_trace( created_at=attempt.created_at, created_by_role=_enum_value(execution.created_by_role), created_by=execution.created_by, - created_by_account=execution.created_by_account, - created_by_end_user=execution.created_by_end_user, + created_by_account=execution.created_by_account(session), + created_by_end_user=execution.created_by_end_user(session), finished_at=attempt.finished_at, inputs_truncated=inputs_truncated, outputs_truncated=outputs_truncated, @@ -173,7 +180,7 @@ def _retry_trace( ) -def _terminal_trace(execution: WorkflowNodeExecutionModel) -> WorkflowNodeExecutionTrace: +def _terminal_trace(execution: WorkflowNodeExecutionModel, *, session: Session) -> WorkflowNodeExecutionTrace: process_data = execution.process_data_dict if process_data is not None and RETRY_HISTORY_PROCESS_DATA_KEY in process_data: process_data = dict(process_data) @@ -202,8 +209,8 @@ def _terminal_trace(execution: WorkflowNodeExecutionModel) -> WorkflowNodeExecut created_at=to_timestamp(execution.created_at), created_by_role=_enum_value(execution.created_by_role), created_by=execution.created_by, - created_by_account=execution.created_by_account, - created_by_end_user=execution.created_by_end_user, + created_by_account=execution.created_by_account(session), + created_by_end_user=execution.created_by_end_user(session), finished_at=to_timestamp(execution.finished_at), inputs_truncated=execution.inputs_truncated, outputs_truncated=execution.outputs_truncated, diff --git a/api/services/workflow_run_service.py b/api/services/workflow_run_service.py index dac000065cd..50c2bc52140 100644 --- a/api/services/workflow_run_service.py +++ b/api/services/workflow_run_service.py @@ -5,6 +5,7 @@ from typing import TypedDict import contexts from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from extensions.ext_database import db from graphon.enums import WorkflowExecutionStatus from libs.infinite_scroll_pagination import InfiniteScrollPagination from machinery.context import RequestContext @@ -202,7 +203,7 @@ class WorkflowRunService: app_id=app_id, workflow_run_id=run_id, ) - return assemble_workflow_node_execution_traces(node_executions, self._node_executions) + return assemble_workflow_node_execution_traces(node_executions, self._node_executions, session=db.session()) def get_pause_details( self, diff --git a/api/tests/test_containers_integration_tests/models/test_workflow_node_execution_model.py b/api/tests/test_containers_integration_tests/models/test_workflow_node_execution_model.py index ecf0af43025..14b8612d657 100644 --- a/api/tests/test_containers_integration_tests/models/test_workflow_node_execution_model.py +++ b/api/tests/test_containers_integration_tests/models/test_workflow_node_execution_model.py @@ -109,7 +109,7 @@ class TestWorkflowNodeExecutionModelCreatedBy: created_by=account.id, ) - result = execution.created_by_account + result = execution.created_by_account(session=db_session_with_containers) assert result is not None assert result.id == account.id @@ -126,7 +126,7 @@ class TestWorkflowNodeExecutionModelCreatedBy: created_by=account.id, ) - result = execution.created_by_account + result = execution.created_by_account(session=db_session_with_containers) assert result is None @@ -146,7 +146,7 @@ class TestWorkflowNodeExecutionModelCreatedBy: created_by=end_user.id, ) - result = execution.created_by_end_user + result = execution.created_by_end_user(session=db_session_with_containers) assert result is not None assert result.id == end_user.id @@ -165,6 +165,6 @@ class TestWorkflowNodeExecutionModelCreatedBy: created_by=end_user.id, ) - result = execution.created_by_end_user + result = execution.created_by_end_user(session=db_session_with_containers) assert result is None diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py index a9686f0cccf..2f9e597bcc1 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py @@ -15,6 +15,7 @@ from sqlalchemy.orm import Session from controllers.common.errors import NotFoundError from controllers.console.app import workflow_run as workflow_run_module from extensions.ext_database import db +from fields.workflow_run_fields import node_execution_response_source from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from machinery.context import RequestContext from models import Account, App, AppMode @@ -348,7 +349,9 @@ def test_workflow_run_node_executions_return_frontend_trace_contract( _account(sqlite_session) execution = _workflow_run_node_execution(sqlite_session) workflow_runs = Mock() - workflow_runs.get_workflow_run_node_executions.return_value = [execution] + workflow_runs.get_workflow_run_node_executions.return_value = [ + node_execution_response_source(execution, session=sqlite_session) + ] _mock_application_services(monkeypatch, workflow_runs) monkeypatch.setattr(db, "session", sqlite_session) request_context = _request_context() diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py index 81b58c32339..46009877a09 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py @@ -45,6 +45,7 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import ( WorkflowListQuery, WorkflowUpdatePayload, ) +from fields.workflow_run_fields import node_execution_response_source from graphon.enums import WorkflowNodeExecutionStatus from libs.datetime_utils import naive_utc_now from models.account import Account, TenantAccountRole @@ -802,8 +803,12 @@ class TestRagPipelineWorkflowRunNodeExecutionListApi: run_id = uuid4() node_exec = make_node_execution(workflow_run_id=str(run_id)) + session_stub = MagicMock() + session_stub.scalar.return_value = None service = MagicMock() - service.get_rag_pipeline_workflow_run_node_executions.return_value = [node_exec] + service.get_rag_pipeline_workflow_run_node_executions.return_value = [ + node_execution_response_source(node_exec, session=session_stub) + ] with ( app.test_request_context("/"), diff --git a/api/tests/unit_tests/models/test_workflow_node_execution_accessors.py b/api/tests/unit_tests/models/test_workflow_node_execution_accessors.py new file mode 100644 index 00000000000..cdb07111dce --- /dev/null +++ b/api/tests/unit_tests/models/test_workflow_node_execution_accessors.py @@ -0,0 +1,110 @@ +"""Regression coverage for ``WorkflowNodeExecutionModel`` account accessors. + +Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch: +``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user`` +looks up an EndUser only when role is END_USER. Also covers the session-carrying response +source used by the console endpoints that validate ``WorkflowRunNodeExecutionResponse`` +with ``from_attributes=True``. +""" + +from uuid import uuid4 + +from sqlalchemy.orm import Session + +from fields.workflow_run_fields import node_execution_response_source +from models.account import Account +from models.enums import CreatorUserRole, EndUserType +from models.model import EndUser +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom + + +def _execution(role: CreatorUserRole, created_by: str) -> WorkflowNodeExecutionModel: + """Construct a WorkflowNodeExecutionModel without touching the database.""" + return WorkflowNodeExecutionModel( + tenant_id="00000000-0000-0000-0000-000000000001", + app_id="00000000-0000-0000-0000-000000000002", + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=None, + index=1, + predecessor_node_id=None, + node_execution_id=None, + node_id="n1", + node_type="start", + title="Start", + inputs=None, + process_data=None, + outputs=None, + status="succeeded", + error=None, + elapsed_time=0.0, + execution_metadata=None, + created_by_role=role, + created_by=created_by, + ) + + +class TestCreatedByAccount: + def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id) + + result = execution.created_by_account(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + execution = _execution(CreatorUserRole.END_USER, created_by=account.id) + + assert execution.created_by_account(session=sqlite_session) is None + + +class TestCreatedByEndUser: + def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + execution = _execution(CreatorUserRole.END_USER, created_by=end_user.id) + + result = execution.created_by_end_user(session=sqlite_session) + + assert result is not None + assert result.id == end_user.id + + def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + execution = _execution(CreatorUserRole.ACCOUNT, created_by=end_user.id) + + assert execution.created_by_end_user(session=sqlite_session) is None + + +class TestNodeExecutionResponseSource: + def test_accessors_resolve_via_wrapped_session_and_other_attributes_proxy(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id) + + source = node_execution_response_source(execution, session=sqlite_session) + + resolved = source.created_by_account + assert resolved is not None + assert resolved.id == account.id + assert source.created_by_end_user is None + assert source.node_id == "n1" diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py index a7194a805ef..258d56c5e02 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py @@ -4,6 +4,7 @@ from collections.abc import Callable from dataclasses import dataclass from datetime import datetime from types import SimpleNamespace +from unittest import mock import pytest from pytest_mock import MockerFixture @@ -1498,7 +1499,7 @@ def test_get_rag_pipeline_workflow_run_node_executions_assembles_configured_repo app_id=pipeline.id, workflow_run_id="run-1", ) - mock_assemble.assert_called_once_with(expected_executions, node_repo) + mock_assemble.assert_called_once_with(expected_executions, node_repo, session=mock.ANY) def test_get_recommended_plugins_returns_empty_when_no_active_plugins( diff --git a/api/tests/unit_tests/services/test_snippet_service.py b/api/tests/unit_tests/services/test_snippet_service.py index 10c1622ab6e..ea2d83c810a 100644 --- a/api/tests/unit_tests/services/test_snippet_service.py +++ b/api/tests/unit_tests/services/test_snippet_service.py @@ -3,7 +3,7 @@ from __future__ import annotations import json from datetime import datetime from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import ANY, Mock import pytest from sqlalchemy import event, select @@ -968,6 +968,8 @@ def test_workflow_run_queries_delegate_to_repositories(monkeypatch: pytest.Monke ) service._workflow_run_repo = workflow_run_repo service._node_execution_service_repo = node_execution_repo + service._session = Mock() + service._session_maker = Mock() snippet = _snippet() expected_traces = [SimpleNamespace(id="node-execution-1:retry:1"), SimpleNamespace(id="node-execution-1")] mock_assemble = Mock(return_value=expected_traces) @@ -998,7 +1000,7 @@ def test_workflow_run_queries_delegate_to_repositories(monkeypatch: pytest.Monke workflow_run_id="run-1", ) mock_assemble.assert_called_once_with( - node_execution_repo.get_executions_by_workflow_run.return_value, node_execution_repo + node_execution_repo.get_executions_by_workflow_run.return_value, node_execution_repo, session=ANY ) node_execution_repo.get_node_last_execution.assert_called_once_with( tenant_id="tenant-1", diff --git a/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py b/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py index a3957c8370d..eb2cc1077f6 100644 --- a/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py +++ b/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py @@ -77,7 +77,7 @@ def _repository(full_process_data: dict[str, object] | None) -> DifyAPIWorkflowN return cast(DifyAPIWorkflowNodeExecutionRepository, repository) -def test_assemble_expands_retry_history_before_terminal_trace() -> None: +def test_assemble_expands_retry_history_before_terminal_trace(sqlite_session: Session) -> None: process_data = { "request": "successful-attempt", RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(2), _retry_attempt(1)], @@ -85,7 +85,7 @@ def test_assemble_expands_retry_history_before_terminal_trace() -> None: execution = _execution(process_data) repository = _repository(process_data) - traces = assemble_workflow_node_execution_traces([execution], repository) + traces = assemble_workflow_node_execution_traces([execution], repository, session=sqlite_session) assert [trace.id for trace in traces] == ["exec-1:retry:1", "exec-1:retry:2", "exec-1"] assert [trace.status for trace in traces] == ["retry", "retry", "succeeded"] @@ -104,18 +104,18 @@ def test_assemble_expands_retry_history_before_terminal_trace() -> None: repository.load_full_process_data.assert_called_once_with(execution) -def test_assemble_keeps_old_execution_without_retry_history() -> None: +def test_assemble_keeps_old_execution_without_retry_history(sqlite_session: Session) -> None: process_data = {"request": "terminal"} execution = _execution(process_data) - traces = assemble_workflow_node_execution_traces([execution], _repository(process_data)) + traces = assemble_workflow_node_execution_traces([execution], _repository(process_data), session=sqlite_session) assert len(traces) == 1 assert traces[0].id == "exec-1" assert traces[0].process_data == process_data -def test_assemble_skips_malformed_and_duplicate_retry_attempts() -> None: +def test_assemble_skips_malformed_and_duplicate_retry_attempts(sqlite_session: Session) -> None: process_data = { RETRY_HISTORY_PROCESS_DATA_KEY: [ _retry_attempt(2), @@ -125,35 +125,41 @@ def test_assemble_skips_malformed_and_duplicate_retry_attempts() -> None: ] } - traces = assemble_workflow_node_execution_traces([_execution(process_data)], _repository(process_data)) + traces = assemble_workflow_node_execution_traces( + [_execution(process_data)], _repository(process_data), session=sqlite_session + ) assert [trace.retry_index for trace in traces[:-1]] == [1, 2] assert traces[0].error == "attempt 1 failed" -def test_assemble_truncates_retry_attempt_fields() -> None: +def test_assemble_truncates_retry_attempt_fields(sqlite_session: Session) -> None: process_data = {RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(1, outputs={"body": "x" * 2_000_000})]} - traces = assemble_workflow_node_execution_traces([_execution(process_data)], _repository(process_data)) + traces = assemble_workflow_node_execution_traces( + [_execution(process_data)], _repository(process_data), session=sqlite_session + ) assert traces[0].outputs_truncated is True assert traces[-1].id == "exec-1" -def test_assemble_falls_back_to_inline_process_data_when_loader_fails() -> None: +def test_assemble_falls_back_to_inline_process_data_when_loader_fails(sqlite_session: Session) -> None: process_data = {RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(1)]} execution = _execution(process_data) repository = _repository(None) repository.load_full_process_data.side_effect = OSError("storage unavailable") - traces = assemble_workflow_node_execution_traces([execution], repository) + traces = assemble_workflow_node_execution_traces([execution], repository, session=sqlite_session) assert [trace.id for trace in traces] == ["exec-1:retry:1", "exec-1"] -def test_virtual_trace_validates_through_node_execution_response() -> None: +def test_virtual_trace_validates_through_node_execution_response(sqlite_session: Session) -> None: process_data = {RETRY_HISTORY_PROCESS_DATA_KEY: [_retry_attempt(1)]} - traces = assemble_workflow_node_execution_traces([_execution(process_data)], _repository(process_data)) + traces = assemble_workflow_node_execution_traces( + [_execution(process_data)], _repository(process_data), session=sqlite_session + ) response = WorkflowRunNodeExecutionListResponse.model_validate({"data": traces}, from_attributes=True) diff --git a/api/tests/unit_tests/services/test_workflow_run_service.py b/api/tests/unit_tests/services/test_workflow_run_service.py index b71829dcf7c..3e7bf683101 100644 --- a/api/tests/unit_tests/services/test_workflow_run_service.py +++ b/api/tests/unit_tests/services/test_workflow_run_service.py @@ -1,7 +1,7 @@ """Unit tests for the Console workflow-run application service.""" from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import ANY, MagicMock import pytest @@ -213,4 +213,4 @@ class TestWorkflowRunServiceQueries: app_id="app-1", workflow_run_id="run-1", ) - mock_assemble.assert_called_once_with(expected_executions, node_executions) + mock_assemble.assert_called_once_with(expected_executions, node_executions, session=ANY)