refactor: use sessionmaker in small services 2 (#34696)

Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/tests/unit_tests/services/test_account_service.py

@@ -1427,16 +1427,18 @@ class TestRegisterService:
         mock_tenant.name = "Test Workspace"
         mock_inviter = TestAccountAssociatedDataFactory.create_account_mock(account_id="inviter-123", name="Inviter")
 
-        # Mock database queries - need to mock the Session query
+        # Mock database queries - need to mock the sessionmaker query
         mock_session = MagicMock()
         mock_session.query.return_value.filter_by.return_value.first.return_value = None  # No existing account
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+        mock_sessionmaker.return_value.begin.return_value.__exit__.return_value = None
+
         with (
-            patch("services.account_service.Session") as mock_session_class,
+            patch("services.account_service.sessionmaker", mock_sessionmaker),
             patch("services.account_service.AccountService.get_account_by_email_with_case_fallback") as mock_lookup,
         ):
-            mock_session_class.return_value.__enter__.return_value = mock_session
-            mock_session_class.return_value.__exit__.return_value = None
             mock_lookup.return_value = None
 
             # Mock RegisterService.register
@@ -1485,12 +1487,14 @@ class TestRegisterService:
         mixed_email = "Invitee@Example.com"
 
         mock_session = MagicMock()
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+        mock_sessionmaker.return_value.begin.return_value.__exit__.return_value = None
+
         with (
-            patch("services.account_service.Session") as mock_session_class,
+            patch("services.account_service.sessionmaker", mock_sessionmaker),
             patch("services.account_service.AccountService.get_account_by_email_with_case_fallback") as mock_lookup,
         ):
-            mock_session_class.return_value.__enter__.return_value = mock_session
-            mock_session_class.return_value.__exit__.return_value = None
             mock_lookup.return_value = None
 
             mock_new_account = TestAccountAssociatedDataFactory.create_account_mock(
@@ -1541,16 +1545,18 @@ class TestRegisterService:
             account_id="existing-user-456", email="existing@example.com", status="pending"
         )
 
-        # Mock database queries - need to mock the Session query
+        # Mock database queries - need to mock the sessionmaker query
         mock_session = MagicMock()
         mock_session.query.return_value.filter_by.return_value.first.return_value = mock_existing_account
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value.__enter__.return_value = mock_session
+        mock_sessionmaker.return_value.begin.return_value.__exit__.return_value = None
+
         with (
-            patch("services.account_service.Session") as mock_session_class,
+            patch("services.account_service.sessionmaker", mock_sessionmaker),
             patch("services.account_service.AccountService.get_account_by_email_with_case_fallback") as mock_lookup,
         ):
-            mock_session_class.return_value.__enter__.return_value = mock_session
-            mock_session_class.return_value.__exit__.return_value = None
             mock_lookup.return_value = mock_existing_account
 
             # Mock scalar for TenantAccountJoin lookup - no existing member

api/tests/unit_tests/services/test_async_workflow_service.py

@@ -357,11 +357,12 @@ class TestAsyncWorkflowService:
         mock_session_context.__enter__.return_value = mock_session
         mock_session_context.__exit__.return_value = None
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = mock_session_context
+
         with (
             patch.object(async_workflow_service_module, "db", new=SimpleNamespace(engine=fake_engine)),
-            patch.object(
-                async_workflow_service_module, "Session", return_value=mock_session_context
-            ) as mock_session_class,
+            patch.object(async_workflow_service_module, "sessionmaker", mock_sessionmaker),
             patch.object(
                 async_workflow_service_module,
                 "SQLAlchemyWorkflowTriggerLogRepository",
@@ -373,7 +374,7 @@ class TestAsyncWorkflowService:
 
         # Assert
         assert result == expected
-        mock_session_class.assert_called_once_with(fake_engine)
+        mock_sessionmaker.assert_called_once_with(fake_engine)
         mock_repo.get_by_id.assert_called_once_with("trigger-log-123", "tenant-123")
 
     def test_should_return_recent_logs_as_dict_list(self):
@@ -391,9 +392,12 @@ class TestAsyncWorkflowService:
         mock_session_context.__enter__.return_value = mock_session
         mock_session_context.__exit__.return_value = None
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = mock_session_context
+
         with (
             patch.object(async_workflow_service_module, "db", new=SimpleNamespace(engine=MagicMock())),
-            patch.object(async_workflow_service_module, "Session", return_value=mock_session_context),
+            patch.object(async_workflow_service_module, "sessionmaker", mock_sessionmaker),
             patch.object(
                 async_workflow_service_module,
                 "SQLAlchemyWorkflowTriggerLogRepository",
@@ -432,9 +436,12 @@ class TestAsyncWorkflowService:
         mock_session_context.__enter__.return_value = mock_session
         mock_session_context.__exit__.return_value = None
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = mock_session_context
+
         with (
             patch.object(async_workflow_service_module, "db", new=SimpleNamespace(engine=MagicMock())),
-            patch.object(async_workflow_service_module, "Session", return_value=mock_session_context),
+            patch.object(async_workflow_service_module, "sessionmaker", mock_sessionmaker),
             patch.object(
                 async_workflow_service_module,
                 "SQLAlchemyWorkflowTriggerLogRepository",

api/tests/unit_tests/services/test_clear_free_plan_tenant_expired_logs.py

@@ -209,8 +209,22 @@ def _session_wrapper_for_no_autoflush(session: Mock) -> Mock:
     return wrapper
 
 
+def _sessionmaker_wrapper_for_begin(session: Mock) -> Mock:
+    """
+    ClearFreePlanTenantExpiredLogs.process uses: with sessionmaker(db.engine).begin() as session:
+    so sessionmaker(db.engine) must return an object with a begin() method that returns a context manager.
+    """
+    begin_cm = MagicMock()
+    begin_cm.__enter__.return_value = session
+    begin_cm.__exit__.return_value = None
+
+    sessionmaker_result = MagicMock()
+    sessionmaker_result.begin.return_value = begin_cm
+    return sessionmaker_result
+
+
 def _session_wrapper_for_direct(session: Mock) -> Mock:
-    """ClearFreePlanTenantExpiredLogs.process uses: with Session(db.engine) as session:"""
+    """ClearFreePlanTenantExpiredLogs.process uses: with Session(db.engine) as session: (for old code paths)"""
     wrapper = MagicMock()
     wrapper.__enter__.return_value = session
     wrapper.__exit__.return_value = None
@@ -348,7 +362,7 @@ def test_process_with_tenant_ids_filters_by_plan_and_logs_errors(monkeypatch: py
     count_query.count.return_value = 2
     count_session.query.return_value = count_query
 
-    monkeypatch.setattr(service_module, "Session", lambda _engine: _session_wrapper_for_direct(count_session))
+    monkeypatch.setattr(service_module, "sessionmaker", lambda _engine: _sessionmaker_wrapper_for_begin(count_session))
 
     # Avoid LocalProxy usage
     flask_app = service_module.Flask("test-app")
@@ -438,8 +452,8 @@ def test_process_without_tenant_ids_batches_and_scales_interval(monkeypatch: pyt
 
     batch_session.query.side_effect = [q1, q2, q3, q4, q_rs]
 
-    sessions = [_session_wrapper_for_direct(total_session), _session_wrapper_for_direct(batch_session)]
-    monkeypatch.setattr(service_module, "Session", lambda _engine: sessions.pop(0))
+    sessions = [_sessionmaker_wrapper_for_begin(total_session), _sessionmaker_wrapper_for_begin(batch_session)]
+    monkeypatch.setattr(service_module, "sessionmaker", lambda _engine: sessions.pop(0))
 
     process_tenant_mock = MagicMock()
     monkeypatch.setattr(ClearFreePlanTenantExpiredLogs, "process_tenant", process_tenant_mock)
@@ -457,7 +471,7 @@ def test_process_with_tenant_ids_emits_progress_every_100(monkeypatch: pytest.Mo
     count_query = MagicMock()
     count_query.count.return_value = 100
     count_session.query.return_value = count_query
-    monkeypatch.setattr(service_module, "Session", lambda _engine: _session_wrapper_for_direct(count_session))
+    monkeypatch.setattr(service_module, "sessionmaker", lambda _engine: _sessionmaker_wrapper_for_begin(count_session))
 
     flask_app = service_module.Flask("test-app")
     monkeypatch.setattr(service_module, "current_app", SimpleNamespace(_get_current_object=lambda: flask_app))
@@ -523,8 +537,8 @@ def test_process_without_tenant_ids_all_intervals_too_many_uses_min_interval(mon
 
     batch_session.query.side_effect = [*count_queries, q_rs]
 
-    sessions = [_session_wrapper_for_direct(total_session), _session_wrapper_for_direct(batch_session)]
-    monkeypatch.setattr(service_module, "Session", lambda _engine: sessions.pop(0))
+    sessions = [_sessionmaker_wrapper_for_begin(total_session), _sessionmaker_wrapper_for_begin(batch_session)]
+    monkeypatch.setattr(service_module, "sessionmaker", lambda _engine: sessions.pop(0))
 
     process_tenant_mock = MagicMock()
     monkeypatch.setattr(ClearFreePlanTenantExpiredLogs, "process_tenant", process_tenant_mock)

api/tests/unit_tests/services/test_dataset_service_dataset.py

@@ -578,26 +578,33 @@ class TestDatasetServiceCreationAndUpdate:
         binding = SimpleNamespace(external_knowledge_id="old-knowledge", external_knowledge_api_id="old-api")
         session = MagicMock()
         session.query.return_value.filter_by.return_value.first.return_value = binding
+        session.add = MagicMock()
         session_context = _make_session_context(session)
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = session_context
+
         with (
             patch("services.dataset_service.db") as mock_db,
-            patch("services.dataset_service.Session", return_value=session_context),
+            patch("services.dataset_service.sessionmaker", mock_sessionmaker),
         ):
             DatasetService._update_external_knowledge_binding("dataset-1", "new-knowledge", "new-api")
 
         assert binding.external_knowledge_id == "new-knowledge"
         assert binding.external_knowledge_api_id == "new-api"
-        mock_db.session.add.assert_called_once_with(binding)
+        session.add.assert_called_once_with(binding)
 
     def test_update_external_knowledge_binding_raises_for_missing_binding(self):
         session = MagicMock()
         session.query.return_value.filter_by.return_value.first.return_value = None
         session_context = _make_session_context(session)
 
+        mock_sessionmaker = MagicMock()
+        mock_sessionmaker.return_value.begin.return_value = session_context
+
         with (
             patch("services.dataset_service.db"),
-            patch("services.dataset_service.Session", return_value=session_context),
+            patch("services.dataset_service.sessionmaker", mock_sessionmaker),
         ):
             with pytest.raises(ValueError, match="External knowledge binding not found"):
                 DatasetService._update_external_knowledge_binding("dataset-1", "knowledge-1", "api-1")