test: migrate core service residuals sessions and ORM models to SQLite

api/tests/unit_tests/conftest.py

@@ -234,20 +234,3 @@ def persist_service_api_dataset_owner(
     """Persist the tenant-owner mapping resolved by dataset-token authentication."""
     session.add_all([tenant, tenant_account_join])
     session.commit()
-
-
-def setup_mock_tenant_owner_execute_result(mock_db: MagicMock, mock_tenant: object, mock_owner: object) -> None:
-    """Stub the legacy owner query; SQLite-backed tests use ``persist_service_api_tenant_owner``."""
-    mock_db.session.execute.return_value.one_or_none.return_value = (mock_tenant, mock_owner)
-
-
-def setup_mock_dataset_owner_execute_result(
-    mock_db: MagicMock,
-    mock_tenant: object,
-    mock_tenant_account_join: object,
-) -> None:
-    """Stub the legacy dataset-owner query; SQLite tests use ``persist_service_api_dataset_owner``."""
-    mock_db.session.execute.return_value.one_or_none.return_value = (
-        mock_tenant,
-        mock_tenant_account_join,
-    )

api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
-from contextlib import contextmanager
+from contextlib import nullcontext
 from datetime import datetime
 from types import SimpleNamespace
 from unittest import mock
 
 import pytest
+from sqlalchemy import select
+from sqlalchemy.orm import Session
 
 from core.app.apps.advanced_chat import generate_task_pipeline as pipeline_module
 from core.app.entities.app_invoke_entities import InvokeFrom
@@ -21,7 +23,7 @@
 from graphon.enums import WorkflowExecutionStatus
 from models.enums import MessageStatus
 from models.execution_extra_content import HumanInputContent
-from models.model import AppMode, EndUser
+from models.model import AppMode, EndUser, Message
 
 
 def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline:
@@ -34,77 +36,58 @@ def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline:
     return pipeline
 
 
-def test_persist_human_input_extra_content_adds_record(monkeypatch: pytest.MonkeyPatch) -> None:
-    pipeline = _build_pipeline()
-    monkeypatch.setattr(pipeline, "_load_human_input_form_id", lambda **kwargs: "form-1")
-
-    captured_session: dict[str, mock.Mock] = {}
+def _message(*, status: MessageStatus, answer: str = "") -> Message:
+    return Message(
+        id="message-1",
+        status=status,
+        answer=answer,
+        invoke_from=InvokeFrom.WEB_APP,
+        from_end_user_id="user-1",
+    )
 
-    @contextmanager
-    def fake_session():
-        session = mock.Mock()
-        session.scalar.return_value = None
-        captured_session["session"] = session
-        yield session
 
-    pipeline._database_session = fake_session  # type: ignore[method-assign]
+def test_persist_human_input_extra_content_adds_record(
+    monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
+) -> None:
+    pipeline = _build_pipeline()
+    monkeypatch.setattr(pipeline, "_load_human_input_form_id", lambda **kwargs: "form-1")
 
     pipeline._persist_human_input_extra_content(node_id="node-1")
 
-    session = captured_session["session"]
-    session.add.assert_called_once()
-    content = session.add.call_args.args[0]
-    assert isinstance(content, HumanInputContent)
+    content = sqlite_session.scalar(select(HumanInputContent))
+    assert content is not None
     assert content.workflow_run_id == "run-1"
     assert content.message_id == "message-1"
     assert content.form_id == "form-1"
 
 
-def test_persist_human_input_extra_content_skips_when_form_missing(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_persist_human_input_extra_content_skips_when_form_missing(
+    monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
+) -> None:
     pipeline = _build_pipeline()
     monkeypatch.setattr(pipeline, "_load_human_input_form_id", lambda **kwargs: None)
 
-    called = {"value": False}
-
-    @contextmanager
-    def fake_session():
-        called["value"] = True
-        session = mock.Mock()
-        yield session
-
-    pipeline._database_session = fake_session  # type: ignore[method-assign]
-
     pipeline._persist_human_input_extra_content(node_id="node-1")
 
-    assert called["value"] is False
+    assert sqlite_session.scalar(select(HumanInputContent)) is None
 
 
-def test_persist_human_input_extra_content_skips_when_existing(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_persist_human_input_extra_content_skips_when_existing(
+    monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
+) -> None:
     pipeline = _build_pipeline()
     monkeypatch.setattr(pipeline, "_load_human_input_form_id", lambda **kwargs: "form-1")
-
-    captured_session: dict[str, mock.Mock] = {}
-
-    @contextmanager
-    def fake_session():
-        session = mock.Mock()
-        session.scalar.return_value = HumanInputContent(
-            workflow_run_id="run-1",
-            message_id="message-1",
-            form_id="form-1",
-        )
-        captured_session["session"] = session
-        yield session
-
-    pipeline._database_session = fake_session  # type: ignore[method-assign]
+    existing = HumanInputContent.new(workflow_run_id="run-1", message_id="message-1", form_id="form-1")
+    sqlite_session.add(existing)
+    sqlite_session.commit()
 
     pipeline._persist_human_input_extra_content(node_id="node-1")
 
-    session = captured_session["session"]
-    session.add.assert_not_called()
+    contents = list(sqlite_session.scalars(select(HumanInputContent)))
+    assert [content.id for content in contents] == [existing.id]
 
 
-def test_handle_workflow_paused_event_persists_human_input_extra_content() -> None:
+def test_handle_workflow_paused_event_persists_human_input_extra_content(unbound_session: Session) -> None:
     pipeline = _build_pipeline()
     pipeline._application_generate_entity = SimpleNamespace(task_id="task-1")
     pipeline._workflow_response_converter = mock.Mock()
@@ -116,19 +99,14 @@ def test_handle_workflow_paused_event_persists_human_input_extra_content() -> No
         ),
     )
     pipeline._save_message = mock.Mock()
-    message = SimpleNamespace(status=MessageStatus.NORMAL)
+    message = _message(status=MessageStatus.NORMAL)
     pipeline._get_message = mock.Mock(return_value=message)
     pipeline._persist_human_input_extra_content = mock.Mock()
     pipeline._base_task_pipeline = mock.Mock()
     pipeline._base_task_pipeline.queue_manager = mock.Mock()
     pipeline._message_saved_on_pause = False
 
-    @contextmanager
-    def fake_session():
-        session = mock.Mock()
-        yield session
-
-    pipeline._database_session = fake_session  # type: ignore[method-assign]
+    pipeline._database_session = lambda: nullcontext(unbound_session)  # type: ignore[method-assign]
 
     reason = HumanInputRequired(
         form_id="form-1",
@@ -147,7 +125,7 @@ def fake_session():
     assert message.status == MessageStatus.PAUSED
 
 
-def test_resume_appends_chunks_to_paused_answer() -> None:
+def test_resume_appends_chunks_to_paused_answer(unbound_session: Session) -> None:
     app_config = SimpleNamespace(app_id="app-1", tenant_id="tenant-1", sensitive_word_avoidance=None)
     application_generate_entity = SimpleNamespace(
         app_config=app_config,
@@ -184,30 +162,12 @@ def test_resume_appends_chunks_to_paused_answer() -> None:
         draft_var_saver_factory=SimpleNamespace(),
     )
 
-    stored_message = SimpleNamespace(
-        id="message-1",
-        answer="before",
-        status=MessageStatus.PAUSED,
-        updated_at=None,
-        provider_response_latency=0,
-        message_tokens=0,
-        message_unit_price=0,
-        message_price_unit=0,
-        answer_tokens=0,
-        answer_unit_price=0,
-        answer_price_unit=0,
-        total_price=0,
-        currency="USD",
-        message_metadata=None,
-        invoke_from=InvokeFrom.WEB_APP,
-        from_account_id=None,
-        from_end_user_id="user-1",
-    )
+    stored_message = _message(status=MessageStatus.PAUSED, answer="before")
     pipeline._get_message = mock.Mock(return_value=stored_message)
     pipeline._recorded_files = []
 
     list(pipeline._handle_text_chunk_event(QueueTextChunkEvent(text="after")))
-    pipeline._save_message(session=mock.Mock())
+    pipeline._save_message(session=unbound_session)
 
     assert stored_message.answer == "beforeafter"
     assert stored_message.status == MessageStatus.NORMAL

api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from contextlib import contextmanager
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import pytest
@@ -493,7 +493,7 @@ def test_workflow_finish_handlers(self):
         pipeline._base_task_pipeline.queue_manager.publish = lambda *args, **kwargs: None
         pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom")
         pipeline._base_task_pipeline.error_to_stream_response = lambda err: err
-        pipeline._get_message = lambda **kwargs: SimpleNamespace(id="message-id")
+        pipeline._get_message = lambda **kwargs: Message(id="message-id")
 
         succeeded_responses = list(pipeline._handle_workflow_succeeded_event(QueueWorkflowSucceededEvent(outputs={})))
         assert len(succeeded_responses) == 2
@@ -587,7 +587,7 @@ def append_new_token(self, text):
         assert result is False
         assert seen == ["token"]
 
-    def test_handle_retriever_and_annotation_events(self, monkeypatch: pytest.MonkeyPatch):
+    def test_handle_retriever_and_annotation_events(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session):
         pipeline = _make_pipeline()
         calls = {"retriever": 0, "annotation": 0}
 
@@ -603,11 +603,7 @@ def _hit_annotation(_manager, event, session):
         retriever_event = QueueRetrieverResourcesEvent(retriever_resources=[])
         annotation_event = QueueAnnotationReplyEvent(message_annotation_id="ann")
 
-        @contextmanager
-        def _fake_session():
-            yield SimpleNamespace()
-
-        monkeypatch.setattr(pipeline, "_database_session", _fake_session)
+        monkeypatch.setattr(pipeline, "_database_session", lambda: nullcontext(unbound_session))
 
         assert list(pipeline._handle_retriever_resources_event(retriever_event)) == []
         assert list(pipeline._handle_annotation_reply_event(annotation_event)) == []

api/tests/unit_tests/core/mcp/test_mcp_client.py

@@ -1,10 +1,12 @@
 """Unit tests for MCP client."""
 
 from contextlib import ExitStack
-from types import TracebackType
+from types import SimpleNamespace, TracebackType
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
+from sqlalchemy import event, text
+from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session
 
 from core.entities.mcp_provider import MCPProviderEntity
@@ -482,14 +484,16 @@ def test_init(self, mock_provider):
         assert client.by_server_id is True
         assert client._has_retried is False
 
-    @patch("core.mcp.auth_client.db")
-    @patch("core.mcp.auth_client.Session")
     @patch("services.tools.mcp_tools_manage_service.MCPToolManageService")
     def test_handle_auth_error_success(
-        self, mock_service_class, mock_session_class, mock_db, auth_client, mock_provider
+        self,
+        mock_service_class,
+        auth_client,
+        mock_provider,
+        monkeypatch: pytest.MonkeyPatch,
+        sqlite_engine: Engine,
     ):
-        mock_session = MagicMock(spec=Session)
-        mock_session_class.return_value.__enter__.return_value = mock_session
+        monkeypatch.setattr("core.mcp.auth_client.db", SimpleNamespace(engine=sqlite_engine))
 
         mock_service = mock_service_class.return_value
         new_provider = MagicMock(spec=MCPProviderEntity)
@@ -507,6 +511,10 @@ def test_handle_auth_error_success(
 
         auth_client._handle_auth_error(error)
 
+        service_session = mock_service_class.call_args.kwargs["session"]
+        assert isinstance(service_session, Session)
+        assert service_session.in_transaction() is False
+
         # Verify service calls - error.resource_metadata_url and error.scope_hint are parsed from header
         mock_service.auth_with_actions.assert_called_once_with(
             mock_provider,
@@ -544,14 +552,17 @@ def test_handle_auth_error_already_retried(self, auth_client):
 
         assert exc_info.value == error
 
-    @patch("core.mcp.auth_client.db")
-    @patch("core.mcp.auth_client.Session")
     @patch("services.tools.mcp_tools_manage_service.MCPToolManageService")
     def test_handle_auth_error_no_token(
-        self, mock_service_class, mock_session_class, mock_db, auth_client, mock_provider
+        self,
+        mock_service_class,
+        auth_client,
+        mock_provider,
+        monkeypatch: pytest.MonkeyPatch,
+        sqlite_engine: Engine,
     ):
         """Test auth error handling when no token is received."""
-        mock_session_class.return_value.__enter__.return_value = MagicMock()
+        monkeypatch.setattr("core.mcp.auth_client.db", SimpleNamespace(engine=sqlite_engine))
         mock_service = mock_service_class.return_value
 
         new_provider = MagicMock(spec=MCPProviderEntity)
@@ -565,28 +576,48 @@ def test_handle_auth_error_no_token(
 
         assert "Authentication failed - no token received" in str(exc_info.value)
 
-    @patch("core.mcp.auth_client.db")
-    @patch("core.mcp.auth_client.Session")
     @patch("services.tools.mcp_tools_manage_service.MCPToolManageService")
-    def test_handle_auth_error_generic_exception(self, mock_service_class, mock_session_class, mock_db, auth_client):
+    def test_handle_auth_error_generic_exception(
+        self,
+        mock_service_class,
+        auth_client,
+        monkeypatch: pytest.MonkeyPatch,
+        sqlite_engine: Engine,
+    ):
         """Test auth error handling when a generic exception occurs."""
-        mock_session_class.side_effect = Exception("DB error")
+        monkeypatch.setattr("core.mcp.auth_client.db", SimpleNamespace(engine=sqlite_engine))
+        service = mock_service_class.return_value
+
+        def fail_statement(*_args, **_kwargs):
+            raise RuntimeError("DB error")
+
+        def execute_statement(*_args, **_kwargs):
+            service_session = mock_service_class.call_args.kwargs["session"]
+            service_session.execute(text("SELECT 1"))
+
+        service.auth_with_actions.side_effect = execute_statement
+        event.listen(sqlite_engine, "before_cursor_execute", fail_statement)
 
         error = MCPAuthError("Auth failed")
 
-        with pytest.raises(MCPAuthError) as exc_info:
-            auth_client._handle_auth_error(error)
+        try:
+            with pytest.raises(MCPAuthError) as exc_info:
+                auth_client._handle_auth_error(error)
+        finally:
+            event.remove(sqlite_engine, "before_cursor_execute", fail_statement)
 
         assert "Authentication retry failed: DB error" in str(exc_info.value)
 
-    @patch("core.mcp.auth_client.db")
-    @patch("core.mcp.auth_client.Session")
     @patch("services.tools.mcp_tools_manage_service.MCPToolManageService")
     def test_handle_auth_error_mcp_auth_error_propagation(
-        self, mock_service_class, mock_session_class, mock_db, auth_client
+        self,
+        mock_service_class,
+        auth_client,
+        monkeypatch: pytest.MonkeyPatch,
+        sqlite_engine: Engine,
     ):
         """Test that MCPAuthError during refresh is propagated as is."""
-        mock_session_class.return_value.__enter__.return_value = MagicMock()
+        monkeypatch.setattr("core.mcp.auth_client.db", SimpleNamespace(engine=sqlite_engine))
         mock_service = mock_service_class.return_value
         mock_service.auth_with_actions.side_effect = MCPAuthError("Refresh failed")