diff --git a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py
index dd95b4231f7..02e8990ebed 100644
--- a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py
+++ b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py
@@ -531,7 +531,7 @@ def test_duplicate_detection_and_id_regeneration(
     sqlite_session_factory: sessionmaker[Session],
 ) -> None:
     repo = _repository(monkeypatch, sqlite_session_factory)
-    duplicate = IntegrityError("duplicate", params=None, orig=Mock(spec=psycopg2.errors.UniqueViolation))
+    duplicate = IntegrityError("duplicate", params=None, orig=psycopg2.errors.UniqueViolation("duplicate key"))
     assert repo._is_duplicate_key_error(duplicate)
     assert not repo._is_duplicate_key_error(IntegrityError("other", params=None, orig=Exception("other")))
     execution = _execution(execution_id="old")
@@ -552,7 +552,7 @@ def test_save_retries_postgres_duplicate_key(
     duplicate = IntegrityError(
         "duplicate",
         params=None,
-        orig=Mock(spec=psycopg2.errors.UniqueViolation),
+        orig=psycopg2.errors.UniqueViolation("duplicate key"),
     )
     persist = Mock(side_effect=[duplicate, None])
     monkeypatch.setattr(repo, "_persist_to_database", persist)
diff --git a/api/tests/unit_tests/core/test_model_manager.py b/api/tests/unit_tests/core/test_model_manager.py
index ca15d835734..cab6013f033 100644
--- a/api/tests/unit_tests/core/test_model_manager.py
+++ b/api/tests/unit_tests/core/test_model_manager.py
@@ -174,7 +174,7 @@ def test_quota_managed_non_streaming_invocation_finalizes_reservation() -> None:
     )
     model_instance = manager.get_model_instance("tenant-1", "openai", ModelType.LLM, "gpt-4")
     usage = LLMUsage.empty_usage().model_copy(update={"total_tokens": 12})
-    result = MagicMock(spec=LLMResult, usage=usage)
+    result = LLMResult(model="gpt-4", prompt_messages=[], message=AssistantPromptMessage(content="answer"), usage=usage)
     reservation = MagicMock(commit_before_delivery=True)
 
     invocation_id = str(uuid4())
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 abfcf2e3437..20104e3c9cb 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
@@ -83,17 +83,20 @@ class MockRepo:
     pass
 
 
-def test_get_published_workflow_by_id_locks_restore_source(mocker: MockerFixture) -> None:
-    session = mocker.Mock(spec=Session)
+def test_get_published_workflow_by_id_locks_restore_source(
+    mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext
+) -> None:
+    session = rag_pipeline_service.session
     workflow = _make_workflow()
     workflow.version = "v1"
-    session.scalar.return_value = workflow
-    service = RagPipelineService.__new__(RagPipelineService)
-    service._session = session
+    session.add(workflow)
+    session.flush()
+    lookup = mocker.spy(session, "scalar")
+    service = rag_pipeline_service.service
 
     result = service.get_published_workflow_by_id(_make_pipeline(), workflow.id)
 
-    stmt = session.scalar.call_args.args[0]
+    stmt = lookup.call_args.args[0]
     sql = str(stmt.compile(dialect=postgresql.dialect()))
     assert result is workflow
     assert "FOR UPDATE" in sql
@@ -1059,11 +1062,26 @@ def test_set_datasource_variables_success(
     draft_wf.get_enclosing_node_type_and_id.return_value = None  # Avoid unpacking error
     mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=draft_wf)
 
-    execution = mocker.Mock(spec=WorkflowNodeExecution)
-    execution.id = "exec-1"
-    execution.process_data = {}
-    execution.inputs = {}
-    execution.outputs = {}
+    execution = WorkflowNodeExecution(
+        id="exec-1",
+        node_execution_id="node-exec-1",
+        workflow_id="wf-1",
+        workflow_execution_id="run-1",
+        index=1,
+        predecessor_node_id=None,
+        node_id="node-1",
+        node_type=BuiltinNodeTypes.DATASOURCE,
+        title="Datasource",
+        process_data={},
+        inputs={},
+        outputs={},
+        status=WorkflowNodeExecutionStatus.SUCCEEDED,
+        error=None,
+        elapsed_time=0,
+        metadata={},
+        created_at=datetime(2026, 1, 1),
+        finished_at=None,
+    )
     mocker.patch.object(rag_pipeline_service.service, "_handle_node_run_result", return_value=execution)
 
     # Mock Repository