test: use sqlite sessions at generation boundaries (#40095)

api/tests/unit_tests/controllers/service_api/app/test_workflow.py

@@ -582,7 +582,7 @@ class TestWorkflowRunApi:
                 handler(api, session=sqlite_session, app_model=app_model, end_user=end_user)
 
     def test_sandbox_billing_does_not_gate_default_workflow_run(
-        self, app: Flask, monkeypatch: pytest.MonkeyPatch
+        self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
     ) -> None:
         workflow_module = sys.modules["controllers.service_api.app.workflow"]
         monkeypatch.setattr(workflow_module.dify_config, "BILLING_ENABLED", True)
@@ -598,7 +598,7 @@ class TestWorkflowRunApi:
         with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}):
             response = handler(
                 api,
-                session=Mock(),
+                session=sqlite_session,
                 app_model=_make_app_model(),
                 end_user=_make_end_user(),
             )
@@ -609,7 +609,9 @@ class TestWorkflowRunApi:
 
 
 class TestWorkflowRunByIdApi:
-    def test_rejects_sandbox_plan_with_upgrade_error(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
+    def test_rejects_sandbox_plan_with_upgrade_error(
+        self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
+    ) -> None:
         workflow_module = sys.modules["controllers.service_api.app.workflow"]
         monkeypatch.setattr(workflow_module.dify_config, "BILLING_ENABLED", True)
 
@@ -626,7 +628,7 @@ class TestWorkflowRunByIdApi:
             with pytest.raises(WorkflowVersionExecutionNotAllowedError) as exc_info:
                 handler(
                     api,
-                    session=Mock(),
+                    session=sqlite_session,
                     app_model=app_model,
                     end_user=_make_end_user(),
                     workflow_id="w1",
@@ -660,6 +662,7 @@ class TestWorkflowRunByIdApi:
         billing_config_enabled: bool,
         billing_enabled: bool,
         plan: CloudPlan,
+        sqlite_session: Session,
     ) -> None:
         workflow_module = sys.modules["controllers.service_api.app.workflow"]
         monkeypatch.setattr(workflow_module.dify_config, "BILLING_ENABLED", billing_config_enabled)
@@ -676,7 +679,7 @@ class TestWorkflowRunByIdApi:
         with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}):
             response = handler(
                 api,
-                session=Mock(),
+                session=sqlite_session,
                 app_model=app_model,
                 end_user=_make_end_user(),
                 workflow_id="w1",

api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py

@@ -10,9 +10,9 @@ from __future__ import annotations
 
 from types import SimpleNamespace
 from typing import Any
-from unittest.mock import MagicMock
 
 import pytest
+from sqlalchemy.orm import Session
 
 import core.app.features.annotation_reply.annotation_reply as annotation_mod
 import core.moderation.input_moderation as input_moderation_mod
@@ -75,13 +75,13 @@ def _saved_user_query(events: list[Any]) -> str:
 
 
 class TestRunInputGuards:
-    def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch):
+    def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session):
         _patch_moderation(monkeypatch, returns=(False, {}, "hello"))
         _patch_annotation(monkeypatch, reply=None)
         qm = _FakeQueueManager()
 
         handled, query, annotation_reply = AgentAppGenerator()._run_input_guards(
-            session=MagicMock(),
+            session=sqlite_session,
             application_generate_entity=_make_entity("hello"),
             app_model=SimpleNamespace(id="app-1"),
             message=SimpleNamespace(id="msg-1"),
@@ -93,13 +93,13 @@ class TestRunInputGuards:
         assert annotation_reply is None
         assert qm.events == []
 
-    def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch):
+    def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session):
         _patch_moderation(monkeypatch, returns=(True, {}, "[redacted]"))
         _patch_annotation(monkeypatch, reply=None)
         qm = _FakeQueueManager()
 
         handled, query, annotation_reply = AgentAppGenerator()._run_input_guards(
-            session=MagicMock(),
+            session=sqlite_session,
             application_generate_entity=_make_entity("leak my secret"),
             app_model=SimpleNamespace(id="app-1"),
             message=SimpleNamespace(id="msg-1"),
@@ -111,13 +111,13 @@ class TestRunInputGuards:
         assert annotation_reply is None
         assert qm.events == []
 
-    def test_moderation_block_short_circuits(self, monkeypatch: pytest.MonkeyPatch):
+    def test_moderation_block_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session):
         _patch_moderation(monkeypatch, raises=ModerationError("blocked preset answer"))
         _patch_annotation(monkeypatch, reply=None)
         qm = _FakeQueueManager()
 
         handled, _, annotation_reply = AgentAppGenerator()._run_input_guards(
-            session=MagicMock(),
+            session=sqlite_session,
             application_generate_entity=_make_entity("forbidden"),
             app_model=SimpleNamespace(id="app-1"),
             message=SimpleNamespace(id="msg-1"),
@@ -130,13 +130,13 @@ class TestRunInputGuards:
         assert _answer_text(qm.events) == "blocked preset answer"
         assert _saved_user_query(qm.events) == "forbidden"
 
-    def test_annotation_hit_short_circuits(self, monkeypatch: pytest.MonkeyPatch):
+    def test_annotation_hit_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session):
         _patch_moderation(monkeypatch, returns=(False, {}, "what is your name"))
         _patch_annotation(monkeypatch, reply=SimpleNamespace(id="anno-1", content="I am the annotated Iris."))
         qm = _FakeQueueManager()
 
         handled, _, annotation_reply = AgentAppGenerator()._run_input_guards(
-            session=MagicMock(),
+            session=sqlite_session,
             application_generate_entity=_make_entity("what is your name"),
             app_model=SimpleNamespace(id="app-1"),
             message=SimpleNamespace(id="msg-1"),

api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py

@@ -5,6 +5,7 @@ import logging
 import pytest
 from pydantic import ValidationError
 from pytest_mock import MockerFixture
+from sqlalchemy.orm import Session, sessionmaker
 
 from core.app.apps.agent_chat.app_generator import AgentChatAppGenerator
 from core.app.apps.exc import GenerateTaskStoppedError
@@ -30,12 +31,12 @@ def generator(mocker: MockerFixture):
 
 
 class TestAgentChatAppGeneratorGenerate:
-    def test_generate_rejects_blocking_mode(self, generator, mocker: MockerFixture):
+    def test_generate_rejects_blocking_mode(self, generator, mocker: MockerFixture, sqlite_session: Session):
         app_model = mocker.MagicMock()
         user = DummyAccount("user")
         with pytest.raises(ValueError):
             generator.generate(
-                session=mocker.MagicMock(),
+                session=sqlite_session,
                 app_model=app_model,
                 user=user,
                 args={},
@@ -43,44 +44,45 @@ class TestAgentChatAppGeneratorGenerate:
                 streaming=False,
             )
 
-    def test_generate_requires_query(self, generator, mocker: MockerFixture):
+    def test_generate_requires_query(self, generator, mocker: MockerFixture, sqlite_session: Session):
         app_model = mocker.MagicMock()
         user = DummyAccount("user")
         with pytest.raises(ValueError):
             generator.generate(
-                session=mocker.MagicMock(),
+                session=sqlite_session,
                 app_model=app_model,
                 user=user,
                 args={"inputs": {}},
                 invoke_from=mocker.MagicMock(),
             )
 
-    def test_generate_rejects_non_string_query(self, generator, mocker: MockerFixture):
+    def test_generate_rejects_non_string_query(self, generator, mocker: MockerFixture, sqlite_session: Session):
         app_model = mocker.MagicMock()
         user = DummyAccount("user")
         with pytest.raises(ValueError):
             generator.generate(
-                session=mocker.MagicMock(),
+                session=sqlite_session,
                 app_model=app_model,
                 user=user,
                 args={"query": 123, "inputs": {}},
                 invoke_from=mocker.MagicMock(),
             )
 
-    def test_generate_override_requires_debugger(self, generator, mocker: MockerFixture):
+    def test_generate_override_requires_debugger(self, generator, mocker: MockerFixture, sqlite_session: Session):
         app_model = mocker.MagicMock()
         user = DummyAccount("user")
+        generator._get_app_model_config = mocker.MagicMock(return_value=mocker.MagicMock())
 
         with pytest.raises(ValueError):
             generator.generate(
-                session=mocker.MagicMock(),
+                session=sqlite_session,
                 app_model=app_model,
                 user=user,
                 args={"query": "hi", "inputs": {}, "model_config": {"model": {"provider": "p"}}},
                 invoke_from=InvokeFrom.WEB_APP,
             )
 
-    def test_generate_success_with_debugger_override(self, generator, mocker: MockerFixture):
+    def test_generate_success_with_debugger_override(self, generator, mocker: MockerFixture, sqlite_session: Session):
         app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent-chat")
         app_model_config = mocker.MagicMock(id="cfg1")
         app_model_config.to_dict.return_value = {"model": {"provider": "p"}}
@@ -154,7 +156,7 @@ class TestAgentChatAppGeneratorGenerate:
             "files": [{"id": "f1"}],
             "trace_session_id": "session-1",
         }
-        session = mocker.MagicMock()
+        session = sqlite_session
 
         result = generator.generate(
             session=session,
@@ -173,7 +175,7 @@ class TestAgentChatAppGeneratorGenerate:
         inspect.signature(worker_call.kwargs["target"]).bind(**worker_call.kwargs["kwargs"])
         thread_obj.start.assert_called_once()
 
-    def test_generate_without_file_config(self, generator, mocker: MockerFixture):
+    def test_generate_without_file_config(self, generator, mocker: MockerFixture, sqlite_session: Session):
         app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent-chat")
         app_model_config = mocker.MagicMock(id="cfg1", app_id="app1")
         app_model_config.to_dict.return_value = {"model": {"provider": "p"}}
@@ -235,7 +237,7 @@ class TestAgentChatAppGeneratorGenerate:
         )
 
         args = {"query": "hello", "inputs": {"name": "world"}}
-        session = mocker.MagicMock()
+        session = sqlite_session
 
         result = generator.generate(
             session=session,
@@ -261,7 +263,9 @@ class TestAgentChatAppGeneratorWorker:
 
         mocker.patch("core.app.apps.agent_chat.app_generator.preserve_flask_contexts", ctx_manager)
 
-    def test_generate_worker_handles_generate_task_stopped(self, generator, mocker: MockerFixture):
+    def test_generate_worker_handles_generate_task_stopped(
+        self, generator, mocker: MockerFixture, sqlite_session_factory: sessionmaker[Session]
+    ):
         queue_manager = mocker.MagicMock()
         generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock())
         generator._get_message = mocker.MagicMock(return_value=mocker.MagicMock())
@@ -269,10 +273,10 @@ class TestAgentChatAppGeneratorWorker:
         runner = mocker.MagicMock()
         runner.run.side_effect = GenerateTaskStoppedError()
         mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner)
-        session_cm = mocker.MagicMock()
-        session_cm.__enter__.return_value = mocker.MagicMock()
-        create_session = mocker.patch("core.app.apps.agent_chat.app_generator.session_factory.create_session")
-        create_session.return_value = session_cm
+        mocker.patch(
+            "core.app.apps.agent_chat.app_generator.session_factory.create_session",
+            side_effect=sqlite_session_factory,
+        )
 
         generator._generate_worker(
             flask_app=mocker.MagicMock(),
@@ -294,7 +298,9 @@ class TestAgentChatAppGeneratorWorker:
             Exception("bad"),
         ],
     )
-    def test_generate_worker_publishes_errors(self, generator, mocker: MockerFixture, error):
+    def test_generate_worker_publishes_errors(
+        self, generator, mocker: MockerFixture, error, sqlite_session_factory: sessionmaker[Session]
+    ):
         queue_manager = mocker.MagicMock()
         generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock())
         generator._get_message = mocker.MagicMock(return_value=mocker.MagicMock())
@@ -302,10 +308,10 @@ class TestAgentChatAppGeneratorWorker:
         runner = mocker.MagicMock()
         runner.run.side_effect = error
         mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner)
-        session_cm = mocker.MagicMock()
-        session_cm.__enter__.return_value = mocker.MagicMock()
-        create_session = mocker.patch("core.app.apps.agent_chat.app_generator.session_factory.create_session")
-        create_session.return_value = session_cm
+        mocker.patch(
+            "core.app.apps.agent_chat.app_generator.session_factory.create_session",
+            side_effect=sqlite_session_factory,
+        )
 
         generator._generate_worker(
             flask_app=mocker.MagicMock(),
@@ -319,7 +325,11 @@ class TestAgentChatAppGeneratorWorker:
         assert queue_manager.publish_error.called
 
     def test_generate_worker_logs_value_error_when_debug(
-        self, generator, mocker: MockerFixture, caplog: pytest.LogCaptureFixture
+        self,
+        generator,
+        mocker: MockerFixture,
+        caplog: pytest.LogCaptureFixture,
+        sqlite_session_factory: sessionmaker[Session],
     ):
         queue_manager = mocker.MagicMock()
         generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock())
@@ -328,10 +338,10 @@ class TestAgentChatAppGeneratorWorker:
         runner = mocker.MagicMock()
         runner.run.side_effect = ValueError("bad")
         mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner)
-        session_cm = mocker.MagicMock()
-        session_cm.__enter__.return_value = mocker.MagicMock()
-        create_session = mocker.patch("core.app.apps.agent_chat.app_generator.session_factory.create_session")
-        create_session.return_value = session_cm
+        mocker.patch(
+            "core.app.apps.agent_chat.app_generator.session_factory.create_session",
+            side_effect=sqlite_session_factory,
+        )
 
         mocker.patch("core.app.apps.agent_chat.app_generator.dify_config", new=mocker.MagicMock(DEBUG=True))
 

api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py

@@ -3,6 +3,7 @@ from typing import cast
 from unittest.mock import MagicMock, patch
 
 import pytest
+from sqlalchemy.orm import Session, sessionmaker
 
 from extensions.logstore.repositories.logstore_workflow_node_execution_repository import (
     LogstoreWorkflowNodeExecutionRepository,
@@ -11,7 +12,9 @@ from models.account import Account
 from models.workflow import WorkflowNodeExecutionTriggeredFrom
 
 
-def test_save_synchronously_writes_sql_when_dual_write_is_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_save_synchronously_writes_sql_when_dual_write_is_disabled(
+    monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]
+) -> None:
     monkeypatch.delenv("LOGSTORE_DUAL_WRITE_ENABLED", raising=False)
     with (
         patch("extensions.logstore.repositories.logstore_workflow_node_execution_repository.AliyunLogStore"),
@@ -21,7 +24,7 @@ def test_save_synchronously_writes_sql_when_dual_write_is_disabled(monkeypatch:
         ) as sql_repository_type,
     ):
         repository = LogstoreWorkflowNodeExecutionRepository(
-            session_factory=MagicMock(),
+            session_factory=sqlite_session_factory,
             tenant_id="tenant-1",
             user=cast(Account, SimpleNamespace(id="account-1")),
             app_id="app-1",

api/tests/unit_tests/services/plugin/test_plugin_migration.py

@@ -1,9 +1,12 @@
 import json
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlalchemy.orm import Session
 
+from models.model import App, AppMode
 from services.plugin.plugin_migration import PluginMigration
 
 MIGRATION_MODULE = "services.plugin.plugin_migration"
@@ -33,24 +36,28 @@ def test_fetch_latest_package_identifier_calls_marketplace_when_enabled(mocker:
     assert result == "langgenius/openai:1.0.0@abc"
 
 
-def test_extract_app_tables_checks_agent_mode_with_its_session(mocker: MockerFixture) -> None:
-    app = mocker.MagicMock(app_model_config_id=None, mode="chat")
-    app.is_agent_with_session.return_value = False
-    apps_result = mocker.MagicMock()
-    apps_result.all.return_value = [app]
-    configs_result = mocker.MagicMock()
-    configs_result.all.return_value = []
-    session = mocker.MagicMock()
-    session.scalars.side_effect = [apps_result, configs_result]
-    session_context = mocker.MagicMock()
-    session_context.__enter__.return_value = session
-    mocker.patch(f"{MIGRATION_MODULE}.Session", return_value=session_context)
-    mocker.patch(f"{MIGRATION_MODULE}.db")
+def test_extract_app_tables_checks_agent_mode_with_its_session(mocker: MockerFixture, sqlite_session: Session) -> None:
+    app = App(
+        id="app-1",
+        tenant_id="tenant-1",
+        name="Chat app",
+        description="",
+        mode=AppMode.CHAT,
+        icon_type=None,
+        icon="",
+        icon_background=None,
+        enable_site=False,
+        enable_api=False,
+        created_by="account-1",
+        max_active_requests=0,
+    )
+    sqlite_session.add(app)
+    sqlite_session.commit()
+    mocker.patch(f"{MIGRATION_MODULE}.db", SimpleNamespace(engine=sqlite_session.get_bind()))
 
     result = PluginMigration.extract_app_tables("tenant-1")
 
     assert result == []
-    app.is_agent_with_session.assert_called_once_with(session=session)
 
 
 class TestHandlePluginInstanceInstall:

api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py

@@ -3,14 +3,51 @@ from typing import cast
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlalchemy.orm import Session
 
 from core.app.entities.app_invoke_entities import InvokeFrom
-from models.dataset import Dataset, Pipeline
+from core.rag.index_processor.constant.index_type import IndexStructureType
+from models.dataset import Dataset, Document, Pipeline
+from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
 from models.model import Account, App, EndUser
 from services.dataset_ref_service import DatasetRefService
 from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService
 
 
+def _make_pipeline(*, tenant_id: str = "tenant-1") -> Pipeline:
+    pipeline = Pipeline(tenant_id=tenant_id, name="Pipeline", description="")
+    pipeline.id = "pipeline-1"
+    return pipeline
+
+
+def _make_dataset(*, dataset_id: str = "dataset-1", tenant_id: str = "tenant-1") -> Dataset:
+    return Dataset(
+        id=dataset_id,
+        tenant_id=tenant_id,
+        name="Dataset",
+        created_by="user-1",
+        pipeline_id="pipeline-1",
+    )
+
+
+def _make_document(
+    *, document_id: str = "doc-1", dataset_id: str = "dataset-1", tenant_id: str = "tenant-1"
+) -> Document:
+    return Document(
+        id=document_id,
+        tenant_id=tenant_id,
+        dataset_id=dataset_id,
+        position=1,
+        data_source_type=DataSourceType.LOCAL_FILE,
+        batch="batch",
+        name="Document",
+        created_from=DocumentCreatedFrom.API,
+        created_by="user-1",
+        indexing_status=IndexingStatus.COMPLETED,
+        doc_form=IndexStructureType.PARAGRAPH_INDEX,
+    )
+
+
 def test_get_max_active_requests_uses_smallest_non_zero_limit(mocker: MockerFixture) -> None:
     mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_DEFAULT_ACTIVE_REQUESTS", 5)
     mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_MAX_ACTIVE_REQUESTS", 3)
@@ -41,14 +78,14 @@ def test_get_max_active_requests_returns_zero_when_all_unlimited(mocker: MockerF
         (InvokeFrom.DEBUGGER, SimpleNamespace(id="wf-1"), None),
     ],
 )
-def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_error) -> None:
+def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_error, sqlite_session: Session) -> None:
     rag_pipeline_service_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.RagPipelineService")
     rag_pipeline_service = rag_pipeline_service_cls.return_value
     rag_pipeline_service.get_draft_workflow.return_value = workflow
     rag_pipeline_service.get_published_workflow.return_value = workflow
 
     pipeline = cast(Pipeline, SimpleNamespace(id="pipeline-1"))
-    session = mocker.Mock()
+    session = sqlite_session
 
     if expected_error:
         with pytest.raises(ValueError, match=expected_error):
@@ -58,7 +95,9 @@ def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_err
         assert result == workflow
 
 
-def test_generate_updates_document_status_and_returns_event_stream(mocker: MockerFixture) -> None:
+def test_generate_updates_document_status_and_returns_event_stream(
+    mocker: MockerFixture, sqlite_session: Session
+) -> None:
     dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1"))
     pipeline = cast(
         Pipeline,
@@ -70,7 +109,6 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke
     )
     user = cast(Account | EndUser, SimpleNamespace(id="user-1"))
     args = {"original_document_id": "doc-1", "query": "hello"}
-    session_mock = mocker.Mock()
 
     mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1"))
     update_status_mock = mocker.patch.object(PipelineGenerateService, "update_document_status")
@@ -86,7 +124,7 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke
         args=args,
         invoke_from=InvokeFrom.WEB_APP,
         streaming=True,
-        session=session_mock,
+        session=sqlite_session,
     )
 
     assert result == "stream-events"
@@ -94,11 +132,11 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke
     assert document_ref.dataset.tenant_id == "tenant-1"
     assert document_ref.dataset.dataset_id == "dataset-1"
     assert document_ref.document_id == "doc-1"
-    update_status_mock.assert_called_once_with(document_ref, session=session_mock)
-    assert generator_instance.generate.call_args.kwargs["session"] is session_mock
+    update_status_mock.assert_called_once_with(document_ref, session=sqlite_session)
+    assert generator_instance.generate.call_args.kwargs["session"] is sqlite_session
 
 
-def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFixture) -> None:
+def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFixture, sqlite_session: Session) -> None:
     dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-2"))
     pipeline = cast(
         Pipeline,
@@ -117,7 +155,7 @@ def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFix
             user=cast(Account, SimpleNamespace(id="user-1")),
             args={"original_document_id": "doc-1"},
             invoke_from=InvokeFrom.WEB_APP,
-            session=mocker.Mock(),
+            session=sqlite_session,
         )
 
     update_status_mock.assert_not_called()
@@ -125,18 +163,13 @@ def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFix
 
 def test_generate_rejects_original_document_outside_pipeline_dataset_before_dispatch(
     mocker: MockerFixture,
+    sqlite_session: Session,
 ) -> None:
-    dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1"))
-    pipeline = cast(
-        Pipeline,
-        SimpleNamespace(
-            id="pipeline-1",
-            tenant_id="tenant-1",
-            retrieve_dataset=mocker.Mock(return_value=dataset),
-        ),
-    )
-    session = mocker.Mock()
-    session.scalar.return_value = None
+    dataset = _make_dataset()
+    pipeline = _make_pipeline()
+    outside_document = _make_document(document_id="foreign-doc", dataset_id="other-dataset", tenant_id="tenant-2")
+    sqlite_session.add_all([dataset, pipeline, outside_document])
+    sqlite_session.commit()
     mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1"))
     generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator")
 
@@ -146,29 +179,25 @@ def test_generate_rejects_original_document_outside_pipeline_dataset_before_disp
             user=cast(Account, SimpleNamespace(id="user-1")),
             args={"original_document_id": "foreign-doc"},
             invoke_from=InvokeFrom.PUBLISHED_PIPELINE,
-            session=session,
+            session=sqlite_session,
         )
 
-    statement = session.scalar.call_args.args[0]
-    assert {"foreign-doc", "dataset-1", "tenant-1"} <= set(statement.compile().params.values())
+    sqlite_session.refresh(outside_document)
+    assert outside_document.indexing_status == IndexingStatus.COMPLETED
     generator_cls.assert_not_called()
 
 
-def test_update_document_status_updates_existing_document(mocker: MockerFixture) -> None:
-    document = SimpleNamespace(indexing_status="completed")
-    dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1"))
+def test_update_document_status_updates_existing_document(sqlite_session: Session) -> None:
+    document = _make_document()
+    dataset = _make_dataset()
+    sqlite_session.add_all([dataset, document])
+    sqlite_session.commit()
     dataset_ref = DatasetRefService.create_dataset_ref(dataset)
     document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1")
 
-    session_mock = mocker.Mock()
-    get_document_mock = mocker.patch.object(DatasetRefService, "get_document_by_ref", return_value=document)
-    add_mock = session_mock.add
+    PipelineGenerateService.update_document_status(document_ref, session=sqlite_session)
 
-    PipelineGenerateService.update_document_status(document_ref, session=session_mock)
-
-    assert document.indexing_status == "waiting"
-    get_document_mock.assert_called_once_with(document_ref, session=session_mock)
-    add_mock.assert_called_once_with(document)
+    assert document.indexing_status == IndexingStatus.WAITING
 
 
 @pytest.mark.parametrize(
@@ -179,42 +208,28 @@ def test_update_document_status_updates_existing_document(mocker: MockerFixture)
     ],
 )
 def test_update_document_status_rejects_document_outside_owner(
-    mocker: MockerFixture,
     document_tenant_id: str,
     document_dataset_id: str,
+    sqlite_session: Session,
 ) -> None:
     dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1"))
     dataset_ref = DatasetRefService.create_dataset_ref(dataset)
     document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1")
-    outside_document = SimpleNamespace(
-        id="doc-1",
-        tenant_id=document_tenant_id,
-        dataset_id=document_dataset_id,
-        indexing_status="completed",
-    )
-    session_mock = mocker.Mock()
-
-    def resolve_document(statement):
-        params = set(statement.compile().params.values())
-        outside_owner = {outside_document.id, outside_document.dataset_id, outside_document.tenant_id}
-        return outside_document if outside_owner <= params else None
-
-    session_mock.scalar.side_effect = resolve_document
-    add_mock = session_mock.add
+    outside_document = _make_document(tenant_id=document_tenant_id, dataset_id=document_dataset_id)
+    sqlite_session.add(outside_document)
+    sqlite_session.commit()
 
     with pytest.raises(ValueError, match="Pipeline document not found"):
-        PipelineGenerateService.update_document_status(document_ref, session=session_mock)
+        PipelineGenerateService.update_document_status(document_ref, session=sqlite_session)
 
-    statement = session_mock.scalar.call_args.args[0]
-    assert {"doc-1", "dataset-1", "tenant-1"} <= set(statement.compile().params.values())
-    assert outside_document.indexing_status == "completed"
-    add_mock.assert_not_called()
+    sqlite_session.refresh(outside_document)
+    assert outside_document.indexing_status == IndexingStatus.COMPLETED
 
 
 # --- generate_single_iteration ---
 
 
-def test_generate_single_iteration_delegates(mocker: MockerFixture) -> None:
+def test_generate_single_iteration_delegates(mocker: MockerFixture, sqlite_session: Session) -> None:
     mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1"))
 
     generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator")
@@ -224,7 +239,7 @@ def test_generate_single_iteration_delegates(mocker: MockerFixture) -> None:
 
     pipeline = cast(Pipeline, SimpleNamespace(id="p1"))
     user = cast(Account, SimpleNamespace(id="u1"))
-    session = mocker.Mock()
+    session = sqlite_session
 
     result = PipelineGenerateService.generate_single_iteration(pipeline, user, "node-1", {"key": "val"}, session)
 
@@ -236,7 +251,7 @@ def test_generate_single_iteration_delegates(mocker: MockerFixture) -> None:
 # --- generate_single_loop ---
 
 
-def test_generate_single_loop_delegates(mocker: MockerFixture) -> None:
+def test_generate_single_loop_delegates(mocker: MockerFixture, sqlite_session: Session) -> None:
     mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1"))
 
     generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator")
@@ -246,7 +261,7 @@ def test_generate_single_loop_delegates(mocker: MockerFixture) -> None:
 
     pipeline = cast(Pipeline, SimpleNamespace(id="p1"))
     user = cast(Account, SimpleNamespace(id="u1"))
-    session = mocker.Mock()
+    session = sqlite_session
 
     result = PipelineGenerateService.generate_single_loop(pipeline, user, "node-1", {"key": "val"}, session)
 

api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py

@@ -362,6 +362,7 @@ class TestBuildWorkflowEventStream:
     def test_build_workflow_event_stream_should_emit_ping_and_terminal_snapshot_event(
         self,
         monkeypatch: pytest.MonkeyPatch,
+        message_session_maker: sessionmaker[Session],
     ) -> None:
         workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.PAUSED)
         topic = _Topic(_StaticSubscription())
@@ -409,7 +410,7 @@ class TestBuildWorkflowEventStream:
                 workflow_run=workflow_run,
                 tenant_id="tenant-1",
                 app_id="app-1",
-                session_maker=MagicMock(),
+                session_maker=message_session_maker,
             )
         )
 
@@ -424,6 +425,7 @@ class TestBuildWorkflowEventStream:
     def test_build_workflow_event_stream_should_emit_periodic_ping_and_stop_after_idle_timeout(
         self,
         monkeypatch: pytest.MonkeyPatch,
+        message_session_maker: sessionmaker[Session],
     ) -> None:
         workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.RUNNING)
         topic = _Topic(_StaticSubscription())
@@ -463,7 +465,7 @@ class TestBuildWorkflowEventStream:
                 workflow_run=workflow_run,
                 tenant_id="tenant-1",
                 app_id="app-1",
-                session_maker=MagicMock(),
+                session_maker=message_session_maker,
                 idle_timeout=20.0,
                 ping_interval=5.0,
             )
@@ -475,6 +477,7 @@ class TestBuildWorkflowEventStream:
     def test_build_workflow_event_stream_should_exit_when_buffer_done_and_empty(
         self,
         monkeypatch: pytest.MonkeyPatch,
+        message_session_maker: sessionmaker[Session],
     ) -> None:
         workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.RUNNING)
         topic = _Topic(_StaticSubscription())
@@ -505,7 +508,7 @@ class TestBuildWorkflowEventStream:
                 workflow_run=workflow_run,
                 tenant_id="tenant-1",
                 app_id="app-1",
-                session_maker=MagicMock(),
+                session_maker=message_session_maker,
             )
         )
 
@@ -515,6 +518,7 @@ class TestBuildWorkflowEventStream:
     def test_build_workflow_event_stream_should_continue_when_pause_loading_fails(
         self,
         monkeypatch: pytest.MonkeyPatch,
+        message_session_maker: sessionmaker[Session],
     ) -> None:
         workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.PAUSED)
         topic = _Topic(_StaticSubscription())
@@ -545,7 +549,7 @@ class TestBuildWorkflowEventStream:
                 workflow_run=workflow_run,
                 tenant_id="tenant-1",
                 app_id="app-1",
-                session_maker=MagicMock(),
+                session_maker=message_session_maker,
             )
         )