refactor(api): use sessionmaker in core app generators & pipelines (#34771)

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

@@ -134,6 +134,7 @@ class TestAdvancedChatAppRunnerConversationVariables:
 
         # Patch the necessary components
         with (
+            patch("core.app.apps.advanced_chat.app_runner.sessionmaker") as mock_sessionmaker,
             patch("core.app.apps.advanced_chat.app_runner.Session") as mock_session_class,
             patch("core.app.apps.advanced_chat.app_runner.select") as mock_select,
             patch("core.app.apps.advanced_chat.app_runner.db") as mock_db,
@@ -150,7 +151,9 @@ class TestAdvancedChatAppRunnerConversationVariables:
             patch("core.app.apps.advanced_chat.app_runner.RedisChannel") as mock_redis_channel_class,
         ):
             # Setup mocks
-            mock_session_class.return_value.__enter__.return_value = mock_session
+            mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+            mock_sessionmaker.return_value.begin.return_value.__exit__ = MagicMock(return_value=False)
+            mock_session_class.return_value.__enter__.return_value = MagicMock()
             mock_db.session.query.return_value.where.return_value.first.return_value = MagicMock()  # App exists
             mock_db.engine = MagicMock()
 
@@ -177,7 +180,6 @@ class TestAdvancedChatAppRunnerConversationVariables:
             # Note: Since we're mocking ConversationVariable.from_variable,
             # we can't directly check the id, but we can verify add_all was called
             assert mock_session.add_all.called, "Session add_all should have been called"
-            assert mock_session.commit.called, "Session commit should have been called"
 
     def test_no_variables_creates_all(self):
         """Test that all conversation variables are created when none exist in DB."""
@@ -278,6 +280,7 @@ class TestAdvancedChatAppRunnerConversationVariables:
 
         # Patch the necessary components
         with (
+            patch("core.app.apps.advanced_chat.app_runner.sessionmaker") as mock_sessionmaker,
             patch("core.app.apps.advanced_chat.app_runner.Session") as mock_session_class,
             patch("core.app.apps.advanced_chat.app_runner.select") as mock_select,
             patch("core.app.apps.advanced_chat.app_runner.db") as mock_db,
@@ -295,7 +298,9 @@ class TestAdvancedChatAppRunnerConversationVariables:
             patch("core.app.apps.advanced_chat.app_runner.RedisChannel") as mock_redis_channel_class,
         ):
             # Setup mocks
-            mock_session_class.return_value.__enter__.return_value = mock_session
+            mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+            mock_sessionmaker.return_value.begin.return_value.__exit__ = MagicMock(return_value=False)
+            mock_session_class.return_value.__enter__.return_value = MagicMock()
             mock_db.session.query.return_value.where.return_value.first.return_value = MagicMock()  # App exists
             mock_db.engine = MagicMock()
 
@@ -326,7 +331,6 @@ class TestAdvancedChatAppRunnerConversationVariables:
             # Verify that all variables were created
             assert len(added_items) == 2, "Should have added both variables"
             assert mock_session.add_all.called, "Session add_all should have been called"
-            assert mock_session.commit.called, "Session commit should have been called"
 
     def test_all_variables_exist_no_changes(self):
         """Test that no changes are made when all variables already exist in DB."""
@@ -429,6 +433,7 @@ class TestAdvancedChatAppRunnerConversationVariables:
 
         # Patch the necessary components
         with (
+            patch("core.app.apps.advanced_chat.app_runner.sessionmaker") as mock_sessionmaker,
             patch("core.app.apps.advanced_chat.app_runner.Session") as mock_session_class,
             patch("core.app.apps.advanced_chat.app_runner.select") as mock_select,
             patch("core.app.apps.advanced_chat.app_runner.db") as mock_db,
@@ -445,7 +450,9 @@ class TestAdvancedChatAppRunnerConversationVariables:
             patch("core.app.apps.advanced_chat.app_runner.RedisChannel") as mock_redis_channel_class,
         ):
             # Setup mocks
-            mock_session_class.return_value.__enter__.return_value = mock_session
+            mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+            mock_sessionmaker.return_value.begin.return_value.__exit__ = MagicMock(return_value=False)
+            mock_session_class.return_value.__enter__.return_value = MagicMock()
             mock_db.session.query.return_value.where.return_value.first.return_value = MagicMock()  # App exists
             mock_db.engine = MagicMock()
 
@@ -465,4 +472,3 @@ class TestAdvancedChatAppRunnerConversationVariables:
 
             # Verify that no variables were added
             assert not mock_session.add_all.called, "Session add_all should not have been called"
-            assert mock_session.commit.called, "Session commit should still be called"

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

@@ -93,6 +93,16 @@ def _patch_common_run_deps(runner: AdvancedChatAppRunner):
                 scalar=lambda *a, **k: MagicMock(),
             ),
         ),
+        sessionmaker=MagicMock(
+            return_value=MagicMock(
+                begin=MagicMock(
+                    return_value=MagicMock(
+                        __enter__=lambda s: MagicMock(scalars=MagicMock(return_value=MagicMock(all=lambda: []))),
+                        __exit__=lambda *a, **k: False,
+                    ),
+                ),
+            ),
+        ),
         select=MagicMock(),
         db=MagicMock(engine=MagicMock()),
         RedisChannel=MagicMock(),

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

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import MagicMock
 
 import pytest
 from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus
@@ -610,33 +611,33 @@ class TestWorkflowGenerateTaskPipeline:
 
     def test_database_session_rolls_back_on_error(self, monkeypatch):
         pipeline = _make_pipeline()
-        calls = {"commit": 0, "rollback": 0}
-
-        class _Session:
-            def __init__(self, *args, **kwargs):
-                _ = args, kwargs
+        calls = {"enter": 0, "exit_exc": None}
 
+        class _BeginContext:
             def __enter__(self):
-                return self
+                calls["enter"] += 1
+                return MagicMock()
 
             def __exit__(self, exc_type, exc, tb):
+                calls["exit_exc"] = exc_type
                 return False
 
-            def commit(self):
-                calls["commit"] += 1
+        class _Sessionmaker:
+            def __init__(self, *args, **kwargs):
+                pass
 
-            def rollback(self):
-                calls["rollback"] += 1
+            def begin(self):
+                return _BeginContext()
 
-        monkeypatch.setattr("core.app.apps.workflow.generate_task_pipeline.Session", _Session)
+        monkeypatch.setattr("core.app.apps.workflow.generate_task_pipeline.sessionmaker", _Sessionmaker)
         monkeypatch.setattr("core.app.apps.workflow.generate_task_pipeline.db", SimpleNamespace(engine=object()))
 
         with pytest.raises(RuntimeError, match="db error"):
             with pipeline._database_session():
                 raise RuntimeError("db error")
 
-        assert calls["commit"] == 0
-        assert calls["rollback"] == 1
+        assert calls["enter"] == 1
+        assert calls["exit_exc"] is RuntimeError
 
     def test_node_retry_and_started_handlers_cover_none_and_value(self):
         pipeline = _make_pipeline()