Merge origin/release/e-1.12.1 into 1.12.1-otel-ee

Sync enterprise 1.12.1 changes:
- feat: implement heartbeat mechanism for database migration lock
- refactor: replace AutoRenewRedisLock with DbMigrationAutoRenewLock
- fix: improve logging for database migration lock release
- fix: make flask upgrade-db fail on error
- fix: include sso_verified in access_mode validation
- fix: inherit web app permission from original app
- fix: make e-1.12.1 enterprise migrations database-agnostic
- fix: get_message_event_type return wrong message type
- refactor: document_indexing_sync_task split db session
- fix: trigger output schema miss
- test: remove unrelated enterprise service test

Conflict resolution:
- Combined OTEL telemetry imports with tool signature import in easy_ui_based_generate_task_pipeline.py

api/tests/integration_tests/tasks/test_remove_app_and_related_data_task.py

@@ -10,7 +10,10 @@ from models import Tenant
 from models.enums import CreatorUserRole
 from models.model import App, UploadFile
 from models.workflow import WorkflowDraftVariable, WorkflowDraftVariableFile
-from tasks.remove_app_and_related_data_task import _delete_draft_variables, delete_draft_variables_batch
+from tasks.remove_app_and_related_data_task import (
+    _delete_draft_variables,
+    delete_draft_variables_batch,
+)
 
 
 @pytest.fixture
@@ -297,12 +300,18 @@ class TestDeleteDraftVariablesWithOffloadIntegration:
     def test_delete_draft_variables_with_offload_data(self, mock_storage, setup_offload_test_data):
         data = setup_offload_test_data
         app_id = data["app"].id
+        upload_file_ids = [uf.id for uf in data["upload_files"]]
+        variable_file_ids = [vf.id for vf in data["variable_files"]]
         mock_storage.delete.return_value = None
 
         with session_factory.create_session() as session:
             draft_vars_before = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
-            var_files_before = session.query(WorkflowDraftVariableFile).count()
-            upload_files_before = session.query(UploadFile).count()
+            var_files_before = (
+                session.query(WorkflowDraftVariableFile)
+                .where(WorkflowDraftVariableFile.id.in_(variable_file_ids))
+                .count()
+            )
+            upload_files_before = session.query(UploadFile).where(UploadFile.id.in_(upload_file_ids)).count()
         assert draft_vars_before == 3
         assert var_files_before == 2
         assert upload_files_before == 2
@@ -315,8 +324,12 @@ class TestDeleteDraftVariablesWithOffloadIntegration:
         assert draft_vars_after == 0
 
         with session_factory.create_session() as session:
-            var_files_after = session.query(WorkflowDraftVariableFile).count()
-            upload_files_after = session.query(UploadFile).count()
+            var_files_after = (
+                session.query(WorkflowDraftVariableFile)
+                .where(WorkflowDraftVariableFile.id.in_(variable_file_ids))
+                .count()
+            )
+            upload_files_after = session.query(UploadFile).where(UploadFile.id.in_(upload_file_ids)).count()
         assert var_files_after == 0
         assert upload_files_after == 0
 
@@ -329,6 +342,8 @@ class TestDeleteDraftVariablesWithOffloadIntegration:
     def test_delete_draft_variables_storage_failure_continues_cleanup(self, mock_storage, setup_offload_test_data):
         data = setup_offload_test_data
         app_id = data["app"].id
+        upload_file_ids = [uf.id for uf in data["upload_files"]]
+        variable_file_ids = [vf.id for vf in data["variable_files"]]
         mock_storage.delete.side_effect = [Exception("Storage error"), None]
 
         deleted_count = delete_draft_variables_batch(app_id, batch_size=10)
@@ -339,8 +354,12 @@ class TestDeleteDraftVariablesWithOffloadIntegration:
         assert draft_vars_after == 0
 
         with session_factory.create_session() as session:
-            var_files_after = session.query(WorkflowDraftVariableFile).count()
-            upload_files_after = session.query(UploadFile).count()
+            var_files_after = (
+                session.query(WorkflowDraftVariableFile)
+                .where(WorkflowDraftVariableFile.id.in_(variable_file_ids))
+                .count()
+            )
+            upload_files_after = session.query(UploadFile).where(UploadFile.id.in_(upload_file_ids)).count()
         assert var_files_after == 0
         assert upload_files_after == 0
 
@@ -395,3 +414,275 @@ class TestDeleteDraftVariablesWithOffloadIntegration:
                 if app2_obj:
                     session.delete(app2_obj)
                 session.commit()
+
+
+class TestDeleteDraftVariablesSessionCommit:
+    """Test suite to verify session commit behavior in delete_draft_variables_batch."""
+
+    @pytest.fixture
+    def setup_offload_test_data(self, app_and_tenant):
+        """Create test data with offload files for session commit tests."""
+        from core.variables.types import SegmentType
+        from libs.datetime_utils import naive_utc_now
+
+        tenant, app = app_and_tenant
+
+        with session_factory.create_session() as session:
+            upload_file1 = UploadFile(
+                tenant_id=tenant.id,
+                storage_type="local",
+                key="test/file1.json",
+                name="file1.json",
+                size=1024,
+                extension="json",
+                mime_type="application/json",
+                created_by_role=CreatorUserRole.ACCOUNT,
+                created_by=str(uuid.uuid4()),
+                created_at=naive_utc_now(),
+                used=False,
+            )
+            upload_file2 = UploadFile(
+                tenant_id=tenant.id,
+                storage_type="local",
+                key="test/file2.json",
+                name="file2.json",
+                size=2048,
+                extension="json",
+                mime_type="application/json",
+                created_by_role=CreatorUserRole.ACCOUNT,
+                created_by=str(uuid.uuid4()),
+                created_at=naive_utc_now(),
+                used=False,
+            )
+            session.add(upload_file1)
+            session.add(upload_file2)
+            session.flush()
+
+            var_file1 = WorkflowDraftVariableFile(
+                tenant_id=tenant.id,
+                app_id=app.id,
+                user_id=str(uuid.uuid4()),
+                upload_file_id=upload_file1.id,
+                size=1024,
+                length=10,
+                value_type=SegmentType.STRING,
+            )
+            var_file2 = WorkflowDraftVariableFile(
+                tenant_id=tenant.id,
+                app_id=app.id,
+                user_id=str(uuid.uuid4()),
+                upload_file_id=upload_file2.id,
+                size=2048,
+                length=20,
+                value_type=SegmentType.OBJECT,
+            )
+            session.add(var_file1)
+            session.add(var_file2)
+            session.flush()
+
+            draft_var1 = WorkflowDraftVariable.new_node_variable(
+                app_id=app.id,
+                node_id="node_1",
+                name="large_var_1",
+                value=StringSegment(value="truncated..."),
+                node_execution_id=str(uuid.uuid4()),
+                file_id=var_file1.id,
+            )
+            draft_var2 = WorkflowDraftVariable.new_node_variable(
+                app_id=app.id,
+                node_id="node_2",
+                name="large_var_2",
+                value=StringSegment(value="truncated..."),
+                node_execution_id=str(uuid.uuid4()),
+                file_id=var_file2.id,
+            )
+            draft_var3 = WorkflowDraftVariable.new_node_variable(
+                app_id=app.id,
+                node_id="node_3",
+                name="regular_var",
+                value=StringSegment(value="regular_value"),
+                node_execution_id=str(uuid.uuid4()),
+            )
+            session.add(draft_var1)
+            session.add(draft_var2)
+            session.add(draft_var3)
+            session.commit()
+
+            data = {
+                "app": app,
+                "tenant": tenant,
+                "upload_files": [upload_file1, upload_file2],
+                "variable_files": [var_file1, var_file2],
+                "draft_variables": [draft_var1, draft_var2, draft_var3],
+            }
+
+        yield data
+
+        with session_factory.create_session() as session:
+            for table, ids in [
+                (WorkflowDraftVariable, [v.id for v in data["draft_variables"]]),
+                (WorkflowDraftVariableFile, [vf.id for vf in data["variable_files"]]),
+                (UploadFile, [uf.id for uf in data["upload_files"]]),
+            ]:
+                cleanup_query = delete(table).where(table.id.in_(ids)).execution_options(synchronize_session=False)
+                session.execute(cleanup_query)
+            session.commit()
+
+    @pytest.fixture
+    def setup_commit_test_data(self, app_and_tenant):
+        """Create test data for session commit tests."""
+        tenant, app = app_and_tenant
+        variable_ids: list[str] = []
+
+        with session_factory.create_session() as session:
+            variables = []
+            for i in range(10):
+                var = WorkflowDraftVariable.new_node_variable(
+                    app_id=app.id,
+                    node_id=f"node_{i}",
+                    name=f"var_{i}",
+                    value=StringSegment(value="test_value"),
+                    node_execution_id=str(uuid.uuid4()),
+                )
+                session.add(var)
+                variables.append(var)
+            session.commit()
+            variable_ids = [v.id for v in variables]
+
+        yield {
+            "app": app,
+            "tenant": tenant,
+            "variable_ids": variable_ids,
+        }
+
+        with session_factory.create_session() as session:
+            cleanup_query = (
+                delete(WorkflowDraftVariable)
+                .where(WorkflowDraftVariable.id.in_(variable_ids))
+                .execution_options(synchronize_session=False)
+            )
+            session.execute(cleanup_query)
+            session.commit()
+
+    def test_session_commit_is_called_after_each_batch(self, setup_commit_test_data):
+        """Test that session.begin() is used for automatic transaction management."""
+        data = setup_commit_test_data
+        app_id = data["app"].id
+
+        # Since session.begin() is used, the transaction is automatically committed
+        # when the with block exits successfully. We verify this by checking that
+        # data is actually persisted.
+        deleted_count = delete_draft_variables_batch(app_id, batch_size=3)
+
+        # Verify all data was deleted (proves transaction was committed)
+        with session_factory.create_session() as session:
+            remaining_count = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
+
+        assert deleted_count == 10
+        assert remaining_count == 0
+
+    def test_data_persisted_after_batch_deletion(self, setup_commit_test_data):
+        """Test that data is actually persisted to database after batch deletion with commits."""
+        data = setup_commit_test_data
+        app_id = data["app"].id
+        variable_ids = data["variable_ids"]
+
+        # Verify initial state
+        with session_factory.create_session() as session:
+            initial_count = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
+        assert initial_count == 10
+
+        # Perform deletion with small batch size to force multiple commits
+        deleted_count = delete_draft_variables_batch(app_id, batch_size=3)
+
+        assert deleted_count == 10
+
+        # Verify all data is deleted in a new session (proves commits worked)
+        with session_factory.create_session() as session:
+            final_count = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
+        assert final_count == 0
+
+        # Verify specific IDs are deleted
+        with session_factory.create_session() as session:
+            remaining_vars = (
+                session.query(WorkflowDraftVariable).where(WorkflowDraftVariable.id.in_(variable_ids)).count()
+            )
+        assert remaining_vars == 0
+
+    def test_session_commit_with_empty_dataset(self, setup_commit_test_data):
+        """Test session behavior when deleting from an empty dataset."""
+        nonexistent_app_id = str(uuid.uuid4())
+
+        # Should not raise any errors and should return 0
+        deleted_count = delete_draft_variables_batch(nonexistent_app_id, batch_size=10)
+        assert deleted_count == 0
+
+    def test_session_commit_with_single_batch(self, setup_commit_test_data):
+        """Test that commit happens correctly when all data fits in a single batch."""
+        data = setup_commit_test_data
+        app_id = data["app"].id
+
+        with session_factory.create_session() as session:
+            initial_count = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
+        assert initial_count == 10
+
+        # Delete all in a single batch
+        deleted_count = delete_draft_variables_batch(app_id, batch_size=100)
+        assert deleted_count == 10
+
+        # Verify data is persisted
+        with session_factory.create_session() as session:
+            final_count = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
+        assert final_count == 0
+
+    def test_invalid_batch_size_raises_error(self, setup_commit_test_data):
+        """Test that invalid batch size raises ValueError."""
+        data = setup_commit_test_data
+        app_id = data["app"].id
+
+        with pytest.raises(ValueError, match="batch_size must be positive"):
+            delete_draft_variables_batch(app_id, batch_size=0)
+
+        with pytest.raises(ValueError, match="batch_size must be positive"):
+            delete_draft_variables_batch(app_id, batch_size=-1)
+
+    @patch("extensions.ext_storage.storage")
+    def test_session_commit_with_offload_data_cleanup(self, mock_storage, setup_offload_test_data):
+        """Test that session commits correctly when cleaning up offload data."""
+        data = setup_offload_test_data
+        app_id = data["app"].id
+        upload_file_ids = [uf.id for uf in data["upload_files"]]
+        mock_storage.delete.return_value = None
+
+        # Verify initial state
+        with session_factory.create_session() as session:
+            draft_vars_before = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
+            var_files_before = (
+                session.query(WorkflowDraftVariableFile)
+                .where(WorkflowDraftVariableFile.id.in_([vf.id for vf in data["variable_files"]]))
+                .count()
+            )
+            upload_files_before = session.query(UploadFile).where(UploadFile.id.in_(upload_file_ids)).count()
+        assert draft_vars_before == 3
+        assert var_files_before == 2
+        assert upload_files_before == 2
+
+        # Delete variables with offload data
+        deleted_count = delete_draft_variables_batch(app_id, batch_size=10)
+        assert deleted_count == 3
+
+        # Verify all data is persisted (deleted) in new session
+        with session_factory.create_session() as session:
+            draft_vars_after = session.query(WorkflowDraftVariable).filter_by(app_id=app_id).count()
+            var_files_after = (
+                session.query(WorkflowDraftVariableFile)
+                .where(WorkflowDraftVariableFile.id.in_([vf.id for vf in data["variable_files"]]))
+                .count()
+            )
+            upload_files_after = session.query(UploadFile).where(UploadFile.id.in_(upload_file_ids)).count()
+        assert draft_vars_after == 0
+        assert var_files_after == 0
+        assert upload_files_after == 0
+
+        # Verify storage cleanup was called
+        assert mock_storage.delete.call_count == 2

api/tests/test_containers_integration_tests/libs/test_auto_renew_redis_lock_integration.py

@@ -0,0 +1,38 @@
+"""
+Integration tests for DbMigrationAutoRenewLock using real Redis via TestContainers.
+"""
+
+import time
+import uuid
+
+import pytest
+
+from extensions.ext_redis import redis_client
+from libs.db_migration_lock import DbMigrationAutoRenewLock
+
+
+@pytest.mark.usefixtures("flask_app_with_containers")
+def test_db_migration_lock_renews_ttl_and_releases():
+    lock_name = f"test:db_migration_auto_renew_lock:{uuid.uuid4().hex}"
+
+    # Keep base TTL very small, and renew frequently so the test is stable even on slower CI.
+    lock = DbMigrationAutoRenewLock(
+        redis_client=redis_client,
+        name=lock_name,
+        ttl_seconds=1.0,
+        renew_interval_seconds=0.2,
+        log_context="test_db_migration_lock",
+    )
+
+    acquired = lock.acquire(blocking=True, blocking_timeout=5)
+    assert acquired is True
+
+    # Wait beyond the base TTL; key should still exist due to renewal.
+    time.sleep(1.5)
+    ttl = redis_client.ttl(lock_name)
+    assert ttl > 0
+
+    lock.release_safely(status="successful")
+
+    # After release, the key should not exist.
+    assert redis_client.exists(lock_name) == 0

api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py

@@ -605,26 +605,20 @@ class TestBatchCreateSegmentToIndexTask:
 
         mock_storage.download.side_effect = mock_download
 
-        # Execute the task
+        # Execute the task - should raise ValueError for empty CSV
         job_id = str(uuid.uuid4())
-        batch_create_segment_to_index_task(
-            job_id=job_id,
-            upload_file_id=upload_file.id,
-            dataset_id=dataset.id,
-            document_id=document.id,
-            tenant_id=tenant.id,
-            user_id=account.id,
-        )
+        with pytest.raises(ValueError, match="The CSV file is empty"):
+            batch_create_segment_to_index_task(
+                job_id=job_id,
+                upload_file_id=upload_file.id,
+                dataset_id=dataset.id,
+                document_id=document.id,
+                tenant_id=tenant.id,
+                user_id=account.id,
+            )
 
         # Verify error handling
-        # Check Redis cache was set to error status
-        from extensions.ext_redis import redis_client
-
-        cache_key = f"segment_batch_import_{job_id}"
-        cache_value = redis_client.get(cache_key)
-        assert cache_value == b"error"
-
-        # Verify no segments were created
+        # Since exception was raised, no segments should be created
         from extensions.ext_database import db
 
         segments = db.session.query(DocumentSegment).all()

api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py

@@ -153,8 +153,7 @@ class TestCleanNotionDocumentTask:
         # Execute cleanup task
         clean_notion_document_task(document_ids, dataset.id)
 
-        # Verify documents and segments are deleted
-        assert db_session_with_containers.query(Document).filter(Document.id.in_(document_ids)).count() == 0
+        # Verify segments are deleted
         assert (
             db_session_with_containers.query(DocumentSegment)
             .filter(DocumentSegment.document_id.in_(document_ids))
@@ -162,9 +161,9 @@ class TestCleanNotionDocumentTask:
             == 0
         )
 
-        # Verify index processor was called for each document
+        # Verify index processor was called
         mock_processor = mock_index_processor_factory.return_value.init_index_processor.return_value
-        assert mock_processor.clean.call_count == len(document_ids)
+        mock_processor.clean.assert_called_once()
 
         # This test successfully verifies:
         # 1. Document records are properly deleted from the database
@@ -186,12 +185,12 @@ class TestCleanNotionDocumentTask:
         non_existent_dataset_id = str(uuid.uuid4())
         document_ids = [str(uuid.uuid4()), str(uuid.uuid4())]
 
-        # Execute cleanup task with non-existent dataset
-        clean_notion_document_task(document_ids, non_existent_dataset_id)
+        # Execute cleanup task with non-existent dataset - expect exception
+        with pytest.raises(Exception, match="Document has no dataset"):
+            clean_notion_document_task(document_ids, non_existent_dataset_id)
 
-        # Verify that the index processor was not called
-        mock_processor = mock_index_processor_factory.return_value.init_index_processor.return_value
-        mock_processor.clean.assert_not_called()
+        # Verify that the index processor factory was not used
+        mock_index_processor_factory.return_value.init_index_processor.assert_not_called()
 
     def test_clean_notion_document_task_empty_document_list(
         self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
@@ -229,9 +228,13 @@ class TestCleanNotionDocumentTask:
         # Execute cleanup task with empty document list
         clean_notion_document_task([], dataset.id)
 
-        # Verify that the index processor was not called
+        # Verify that the index processor was called once with empty node list
         mock_processor = mock_index_processor_factory.return_value.init_index_processor.return_value
-        mock_processor.clean.assert_not_called()
+        assert mock_processor.clean.call_count == 1
+        args, kwargs = mock_processor.clean.call_args
+        # args: (dataset, total_index_node_ids)
+        assert isinstance(args[0], Dataset)
+        assert args[1] == []
 
     def test_clean_notion_document_task_with_different_index_types(
         self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
@@ -315,8 +318,7 @@ class TestCleanNotionDocumentTask:
             # Note: This test successfully verifies cleanup with different document types.
             # The task properly handles various index types and document configurations.
 
-            # Verify documents and segments are deleted
-            assert db_session_with_containers.query(Document).filter(Document.id == document.id).count() == 0
+            # Verify segments are deleted
             assert (
                 db_session_with_containers.query(DocumentSegment)
                 .filter(DocumentSegment.document_id == document.id)
@@ -404,8 +406,7 @@ class TestCleanNotionDocumentTask:
         # Execute cleanup task
         clean_notion_document_task([document.id], dataset.id)
 
-        # Verify documents and segments are deleted
-        assert db_session_with_containers.query(Document).filter(Document.id == document.id).count() == 0
+        # Verify segments are deleted
         assert (
             db_session_with_containers.query(DocumentSegment).filter(DocumentSegment.document_id == document.id).count()
             == 0
@@ -508,8 +509,7 @@ class TestCleanNotionDocumentTask:
 
         clean_notion_document_task(documents_to_clean, dataset.id)
 
-        # Verify only specified documents and segments are deleted
-        assert db_session_with_containers.query(Document).filter(Document.id.in_(documents_to_clean)).count() == 0
+        # Verify only specified documents' segments are deleted
         assert (
             db_session_with_containers.query(DocumentSegment)
             .filter(DocumentSegment.document_id.in_(documents_to_clean))
@@ -697,11 +697,12 @@ class TestCleanNotionDocumentTask:
         db_session_with_containers.commit()
 
         # Mock index processor to raise an exception
-        mock_index_processor = mock_index_processor_factory.init_index_processor.return_value
+        mock_index_processor = mock_index_processor_factory.return_value.init_index_processor.return_value
         mock_index_processor.clean.side_effect = Exception("Index processor error")
 
-        # Execute cleanup task - it should handle the exception gracefully
-        clean_notion_document_task([document.id], dataset.id)
+        # Execute cleanup task - current implementation propagates the exception
+        with pytest.raises(Exception, match="Index processor error"):
+            clean_notion_document_task([document.id], dataset.id)
 
         # Note: This test demonstrates the task's error handling capability.
         # Even with external service errors, the database operations complete successfully.
@@ -803,8 +804,7 @@ class TestCleanNotionDocumentTask:
         all_document_ids = [doc.id for doc in documents]
         clean_notion_document_task(all_document_ids, dataset.id)
 
-        # Verify all documents and segments are deleted
-        assert db_session_with_containers.query(Document).filter(Document.dataset_id == dataset.id).count() == 0
+        # Verify all segments are deleted
         assert (
             db_session_with_containers.query(DocumentSegment).filter(DocumentSegment.dataset_id == dataset.id).count()
             == 0
@@ -914,8 +914,7 @@ class TestCleanNotionDocumentTask:
 
         clean_notion_document_task([target_document.id], target_dataset.id)
 
-        # Verify only documents from target dataset are deleted
-        assert db_session_with_containers.query(Document).filter(Document.id == target_document.id).count() == 0
+        # Verify only documents' segments from target dataset are deleted
         assert (
             db_session_with_containers.query(DocumentSegment)
             .filter(DocumentSegment.document_id == target_document.id)
@@ -1030,8 +1029,7 @@ class TestCleanNotionDocumentTask:
         all_document_ids = [doc.id for doc in documents]
         clean_notion_document_task(all_document_ids, dataset.id)
 
-        # Verify all documents and segments are deleted regardless of status
-        assert db_session_with_containers.query(Document).filter(Document.dataset_id == dataset.id).count() == 0
+        # Verify all segments are deleted regardless of status
         assert (
             db_session_with_containers.query(DocumentSegment).filter(DocumentSegment.dataset_id == dataset.id).count()
             == 0
@@ -1142,8 +1140,7 @@ class TestCleanNotionDocumentTask:
         # Execute cleanup task
         clean_notion_document_task([document.id], dataset.id)
 
-        # Verify documents and segments are deleted
-        assert db_session_with_containers.query(Document).filter(Document.id == document.id).count() == 0
+        # Verify segments are deleted
         assert (
             db_session_with_containers.query(DocumentSegment).filter(DocumentSegment.document_id == document.id).count()
             == 0

api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py

@@ -0,0 +1,182 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from faker import Faker
+
+from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
+from models.dataset import Dataset, Document, DocumentSegment
+from tasks.document_indexing_update_task import document_indexing_update_task
+
+
+class TestDocumentIndexingUpdateTask:
+    @pytest.fixture
+    def mock_external_dependencies(self):
+        """Patch external collaborators used by the update task.
+        - IndexProcessorFactory.init_index_processor().clean(...)
+        - IndexingRunner.run([...])
+        """
+        with (
+            patch("tasks.document_indexing_update_task.IndexProcessorFactory") as mock_factory,
+            patch("tasks.document_indexing_update_task.IndexingRunner") as mock_runner,
+        ):
+            processor_instance = MagicMock()
+            mock_factory.return_value.init_index_processor.return_value = processor_instance
+
+            runner_instance = MagicMock()
+            mock_runner.return_value = runner_instance
+
+            yield {
+                "factory": mock_factory,
+                "processor": processor_instance,
+                "runner": mock_runner,
+                "runner_instance": runner_instance,
+            }
+
+    def _create_dataset_document_with_segments(self, db_session_with_containers, *, segment_count: int = 2):
+        fake = Faker()
+
+        # Account and tenant
+        account = Account(
+            email=fake.email(),
+            name=fake.name(),
+            interface_language="en-US",
+            status="active",
+        )
+        db_session_with_containers.add(account)
+        db_session_with_containers.commit()
+
+        tenant = Tenant(name=fake.company(), status="normal")
+        db_session_with_containers.add(tenant)
+        db_session_with_containers.commit()
+
+        join = TenantAccountJoin(
+            tenant_id=tenant.id,
+            account_id=account.id,
+            role=TenantAccountRole.OWNER,
+            current=True,
+        )
+        db_session_with_containers.add(join)
+        db_session_with_containers.commit()
+
+        # Dataset and document
+        dataset = Dataset(
+            tenant_id=tenant.id,
+            name=fake.company(),
+            description=fake.text(max_nb_chars=64),
+            data_source_type="upload_file",
+            indexing_technique="high_quality",
+            created_by=account.id,
+        )
+        db_session_with_containers.add(dataset)
+        db_session_with_containers.commit()
+
+        document = Document(
+            tenant_id=tenant.id,
+            dataset_id=dataset.id,
+            position=0,
+            data_source_type="upload_file",
+            batch="test_batch",
+            name=fake.file_name(),
+            created_from="upload_file",
+            created_by=account.id,
+            indexing_status="waiting",
+            enabled=True,
+            doc_form="text_model",
+        )
+        db_session_with_containers.add(document)
+        db_session_with_containers.commit()
+
+        # Segments
+        node_ids = []
+        for i in range(segment_count):
+            node_id = f"node-{i + 1}"
+            seg = DocumentSegment(
+                tenant_id=tenant.id,
+                dataset_id=dataset.id,
+                document_id=document.id,
+                position=i,
+                content=fake.text(max_nb_chars=32),
+                answer=None,
+                word_count=10,
+                tokens=5,
+                index_node_id=node_id,
+                status="completed",
+                created_by=account.id,
+            )
+            db_session_with_containers.add(seg)
+            node_ids.append(node_id)
+        db_session_with_containers.commit()
+
+        # Refresh to ensure ORM state
+        db_session_with_containers.refresh(dataset)
+        db_session_with_containers.refresh(document)
+
+        return dataset, document, node_ids
+
+    def test_cleans_segments_and_reindexes(self, db_session_with_containers, mock_external_dependencies):
+        dataset, document, node_ids = self._create_dataset_document_with_segments(db_session_with_containers)
+
+        # Act
+        document_indexing_update_task(dataset.id, document.id)
+
+        # Ensure we see committed changes from another session
+        db_session_with_containers.expire_all()
+
+        # Assert document status updated before reindex
+        updated = db_session_with_containers.query(Document).where(Document.id == document.id).first()
+        assert updated.indexing_status == "parsing"
+        assert updated.processing_started_at is not None
+
+        # Segments should be deleted
+        remaining = (
+            db_session_with_containers.query(DocumentSegment).where(DocumentSegment.document_id == document.id).count()
+        )
+        assert remaining == 0
+
+        # Assert index processor clean was called with expected args
+        clean_call = mock_external_dependencies["processor"].clean.call_args
+        assert clean_call is not None
+        args, kwargs = clean_call
+        # args[0] is a Dataset instance (from another session) — validate by id
+        assert getattr(args[0], "id", None) == dataset.id
+        # args[1] should contain our node_ids
+        assert set(args[1]) == set(node_ids)
+        assert kwargs.get("with_keywords") is True
+        assert kwargs.get("delete_child_chunks") is True
+
+        # Assert indexing runner invoked with the updated document
+        run_call = mock_external_dependencies["runner_instance"].run.call_args
+        assert run_call is not None
+        run_docs = run_call[0][0]
+        assert len(run_docs) == 1
+        first = run_docs[0]
+        assert getattr(first, "id", None) == document.id
+
+    def test_clean_error_is_logged_and_indexing_continues(self, db_session_with_containers, mock_external_dependencies):
+        dataset, document, node_ids = self._create_dataset_document_with_segments(db_session_with_containers)
+
+        # Force clean to raise; task should continue to indexing
+        mock_external_dependencies["processor"].clean.side_effect = Exception("boom")
+
+        document_indexing_update_task(dataset.id, document.id)
+
+        # Ensure we see committed changes from another session
+        db_session_with_containers.expire_all()
+
+        # Indexing should still be triggered
+        mock_external_dependencies["runner_instance"].run.assert_called_once()
+
+        # Segments should remain (since clean failed before DB delete)
+        remaining = (
+            db_session_with_containers.query(DocumentSegment).where(DocumentSegment.document_id == document.id).count()
+        )
+        assert remaining > 0
+
+    def test_document_not_found_noop(self, db_session_with_containers, mock_external_dependencies):
+        fake = Faker()
+        # Act with non-existent document id
+        document_indexing_update_task(dataset_id=fake.uuid4(), document_id=fake.uuid4())
+
+        # Neither processor nor runner should be called
+        mock_external_dependencies["processor"].clean.assert_not_called()
+        mock_external_dependencies["runner_instance"].run.assert_not_called()

api/tests/unit_tests/commands/test_upgrade_db.py

@@ -0,0 +1,146 @@
+import sys
+import threading
+import types
+from unittest.mock import MagicMock
+
+import commands
+from libs.db_migration_lock import LockNotOwnedError, RedisError
+
+HEARTBEAT_WAIT_TIMEOUT_SECONDS = 5.0
+
+
+def _install_fake_flask_migrate(monkeypatch, upgrade_impl) -> None:
+    module = types.ModuleType("flask_migrate")
+    module.upgrade = upgrade_impl
+    monkeypatch.setitem(sys.modules, "flask_migrate", module)
+
+
+def _invoke_upgrade_db() -> int:
+    try:
+        commands.upgrade_db.callback()
+    except SystemExit as e:
+        return int(e.code or 0)
+    return 0
+
+
+def test_upgrade_db_skips_when_lock_not_acquired(monkeypatch, capsys):
+    monkeypatch.setattr(commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 1234)
+
+    lock = MagicMock()
+    lock.acquire.return_value = False
+    commands.redis_client.lock.return_value = lock
+
+    exit_code = _invoke_upgrade_db()
+    captured = capsys.readouterr()
+
+    assert exit_code == 0
+    assert "Database migration skipped" in captured.out
+
+    commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=1234, thread_local=False)
+    lock.acquire.assert_called_once_with(blocking=False)
+    lock.release.assert_not_called()
+
+
+def test_upgrade_db_failure_not_masked_by_lock_release(monkeypatch, capsys):
+    monkeypatch.setattr(commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 321)
+
+    lock = MagicMock()
+    lock.acquire.return_value = True
+    lock.release.side_effect = LockNotOwnedError("simulated")
+    commands.redis_client.lock.return_value = lock
+
+    def _upgrade():
+        raise RuntimeError("boom")
+
+    _install_fake_flask_migrate(monkeypatch, _upgrade)
+
+    exit_code = _invoke_upgrade_db()
+    captured = capsys.readouterr()
+
+    assert exit_code == 1
+    assert "Database migration failed: boom" in captured.out
+
+    commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=321, thread_local=False)
+    lock.acquire.assert_called_once_with(blocking=False)
+    lock.release.assert_called_once()
+
+
+def test_upgrade_db_success_ignores_lock_not_owned_on_release(monkeypatch, capsys):
+    monkeypatch.setattr(commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 999)
+
+    lock = MagicMock()
+    lock.acquire.return_value = True
+    lock.release.side_effect = LockNotOwnedError("simulated")
+    commands.redis_client.lock.return_value = lock
+
+    _install_fake_flask_migrate(monkeypatch, lambda: None)
+
+    exit_code = _invoke_upgrade_db()
+    captured = capsys.readouterr()
+
+    assert exit_code == 0
+    assert "Database migration successful!" in captured.out
+
+    commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=999, thread_local=False)
+    lock.acquire.assert_called_once_with(blocking=False)
+    lock.release.assert_called_once()
+
+
+def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys):
+    """
+    Ensure the lock is renewed while migrations are running, so the base TTL can stay short.
+    """
+
+    # Use a small TTL so the heartbeat interval triggers quickly.
+    monkeypatch.setattr(commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3)
+
+    lock = MagicMock()
+    lock.acquire.return_value = True
+    commands.redis_client.lock.return_value = lock
+
+    renewed = threading.Event()
+
+    def _reacquire():
+        renewed.set()
+        return True
+
+    lock.reacquire.side_effect = _reacquire
+
+    def _upgrade():
+        assert renewed.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS)
+
+    _install_fake_flask_migrate(monkeypatch, _upgrade)
+
+    exit_code = _invoke_upgrade_db()
+    _ = capsys.readouterr()
+
+    assert exit_code == 0
+    assert lock.reacquire.call_count >= 1
+
+
+def test_upgrade_db_ignores_reacquire_errors(monkeypatch, capsys):
+    # Use a small TTL so heartbeat runs during the upgrade call.
+    monkeypatch.setattr(commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3)
+
+    lock = MagicMock()
+    lock.acquire.return_value = True
+    commands.redis_client.lock.return_value = lock
+
+    attempted = threading.Event()
+
+    def _reacquire():
+        attempted.set()
+        raise RedisError("simulated")
+
+    lock.reacquire.side_effect = _reacquire
+
+    def _upgrade():
+        assert attempted.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS)
+
+    _install_fake_flask_migrate(monkeypatch, _upgrade)
+
+    exit_code = _invoke_upgrade_db()
+    _ = capsys.readouterr()
+
+    assert exit_code == 0
+    assert lock.reacquire.call_count >= 1

api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py

@@ -25,15 +25,19 @@ class TestMessageCycleManagerOptimization:
         task_state = Mock()
         return MessageCycleManager(application_generate_entity=mock_application_generate_entity, task_state=task_state)
 
-    def test_get_message_event_type_with_message_file(self, message_cycle_manager):
-        """Test get_message_event_type returns MESSAGE_FILE when message has files."""
+    def test_get_message_event_type_with_assistant_file(self, message_cycle_manager):
+        """Test get_message_event_type returns MESSAGE_FILE when message has assistant-generated files.
+
+        This ensures that AI-generated images (belongs_to='assistant') trigger the MESSAGE_FILE event,
+        allowing the frontend to properly display generated image files with url field.
+        """
         with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
             # Setup mock session and message file
             mock_session = Mock()
             mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
 
             mock_message_file = Mock()
-            # Current implementation uses session.scalar(select(...))
+            mock_message_file.belongs_to = "assistant"
             mock_session.scalar.return_value = mock_message_file
 
             # Execute
@@ -44,6 +48,31 @@ class TestMessageCycleManagerOptimization:
             assert result == StreamEvent.MESSAGE_FILE
             mock_session.scalar.assert_called_once()
 
+    def test_get_message_event_type_with_user_file(self, message_cycle_manager):
+        """Test get_message_event_type returns MESSAGE when message only has user-uploaded files.
+
+        This is a regression test for the issue where user-uploaded images (belongs_to='user')
+        caused the LLM text response to be incorrectly tagged with MESSAGE_FILE event,
+        resulting in broken images in the chat UI. The query filters for belongs_to='assistant',
+        so when only user files exist, the database query returns None, resulting in MESSAGE event type.
+        """
+        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
+            # Setup mock session and message file
+            mock_session = Mock()
+            mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
+
+            # When querying for assistant files with only user files present, return None
+            # (simulates database query with belongs_to='assistant' filter returning no results)
+            mock_session.scalar.return_value = None
+
+            # Execute
+            with current_app.app_context():
+                result = message_cycle_manager.get_message_event_type("test-message-id")
+
+            # Assert
+            assert result == StreamEvent.MESSAGE
+            mock_session.scalar.assert_called_once()
+
     def test_get_message_event_type_without_message_file(self, message_cycle_manager):
         """Test get_message_event_type returns MESSAGE when message has no files."""
         with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
@@ -69,7 +98,7 @@ class TestMessageCycleManagerOptimization:
             mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
 
             mock_message_file = Mock()
-            # Current implementation uses session.scalar(select(...))
+            mock_message_file.belongs_to = "assistant"
             mock_session.scalar.return_value = mock_message_file
 
             # Execute: compute event type once, then pass to message_to_stream_response

api/tests/unit_tests/factories/test_variable_factory.py

@@ -4,7 +4,7 @@ from typing import Any
 from uuid import uuid4
 
 import pytest
-from hypothesis import given, settings
+from hypothesis import HealthCheck, given, settings
 from hypothesis import strategies as st
 
 from core.file import File, FileTransferMethod, FileType
@@ -493,7 +493,7 @@ def _scalar_value() -> st.SearchStrategy[int | float | str | File | None]:
     )
 
 
-@settings(max_examples=50)
+@settings(max_examples=30, suppress_health_check=[HealthCheck.too_slow, HealthCheck.filter_too_much], deadline=None)
 @given(_scalar_value())
 def test_build_segment_and_extract_values_for_scalar_types(value):
     seg = variable_factory.build_segment(value)
@@ -504,7 +504,7 @@ def test_build_segment_and_extract_values_for_scalar_types(value):
         assert seg.value == value
 
 
-@settings(max_examples=50)
+@settings(max_examples=30, suppress_health_check=[HealthCheck.too_slow, HealthCheck.filter_too_much], deadline=None)
 @given(values=st.lists(_scalar_value(), max_size=20))
 def test_build_segment_and_extract_values_for_array_types(values):
     seg = variable_factory.build_segment(values)

api/tests/unit_tests/tasks/test_dataset_indexing_task.py

@@ -83,23 +83,127 @@ def mock_documents(document_ids, dataset_id):
 def mock_db_session():
     """Mock database session via session_factory.create_session()."""
     with patch("tasks.document_indexing_task.session_factory") as mock_sf:
-        session = MagicMock()
-        # Ensure tests that expect session.close() to be called can observe it via the context manager
-        session.close = MagicMock()
-        cm = MagicMock()
-        cm.__enter__.return_value = session
-        # Link __exit__ to session.close so "close" expectations reflect context manager teardown
+        sessions = []  # Track all created sessions
+        # Shared mock data that all sessions will access
+        shared_mock_data = {"dataset": None, "documents": None, "doc_iter": None}
 
-        def _exit_side_effect(*args, **kwargs):
-            session.close()
+        def create_session_side_effect():
+            session = MagicMock()
+            session.close = MagicMock()
 
-        cm.__exit__.side_effect = _exit_side_effect
-        mock_sf.create_session.return_value = cm
+            # Track commit calls
+            commit_mock = MagicMock()
+            session.commit = commit_mock
+            cm = MagicMock()
+            cm.__enter__.return_value = session
 
-        query = MagicMock()
-        session.query.return_value = query
-        query.where.return_value = query
-        yield session
+            def _exit_side_effect(*args, **kwargs):
+                session.close()
+
+            cm.__exit__.side_effect = _exit_side_effect
+
+            # Support session.begin() for transactions
+            begin_cm = MagicMock()
+            begin_cm.__enter__.return_value = session
+
+            def begin_exit_side_effect(*args, **kwargs):
+                # Auto-commit on transaction exit (like SQLAlchemy)
+                session.commit()
+                # Also mark wrapper's commit as called
+                if sessions:
+                    sessions[0].commit()
+
+            begin_cm.__exit__ = MagicMock(side_effect=begin_exit_side_effect)
+            session.begin = MagicMock(return_value=begin_cm)
+
+            sessions.append(session)
+
+            # Setup query with side_effect to handle both Dataset and Document queries
+            def query_side_effect(*args):
+                query = MagicMock()
+                if args and args[0] == Dataset and shared_mock_data["dataset"] is not None:
+                    where_result = MagicMock()
+                    where_result.first.return_value = shared_mock_data["dataset"]
+                    query.where = MagicMock(return_value=where_result)
+                elif args and args[0] == Document and shared_mock_data["documents"] is not None:
+                    # Support both .first() and .all() calls with chaining
+                    where_result = MagicMock()
+                    where_result.where = MagicMock(return_value=where_result)
+
+                    # Create an iterator for .first() calls if not exists
+                    if shared_mock_data["doc_iter"] is None:
+                        docs = shared_mock_data["documents"] or [None]
+                        shared_mock_data["doc_iter"] = iter(docs)
+
+                    where_result.first = lambda: next(shared_mock_data["doc_iter"], None)
+                    docs_or_empty = shared_mock_data["documents"] or []
+                    where_result.all = MagicMock(return_value=docs_or_empty)
+                    query.where = MagicMock(return_value=where_result)
+                else:
+                    query.where = MagicMock(return_value=query)
+                return query
+
+            session.query = MagicMock(side_effect=query_side_effect)
+            return cm
+
+        mock_sf.create_session.side_effect = create_session_side_effect
+
+        # Create a wrapper that behaves like the first session but has access to all sessions
+        class SessionWrapper:
+            def __init__(self):
+                self._sessions = sessions
+                self._shared_data = shared_mock_data
+                # Create a default session for setup phase
+                self._default_session = MagicMock()
+                self._default_session.close = MagicMock()
+                self._default_session.commit = MagicMock()
+
+                # Support session.begin() for default session too
+                begin_cm = MagicMock()
+                begin_cm.__enter__.return_value = self._default_session
+
+                def default_begin_exit_side_effect(*args, **kwargs):
+                    self._default_session.commit()
+
+                begin_cm.__exit__ = MagicMock(side_effect=default_begin_exit_side_effect)
+                self._default_session.begin = MagicMock(return_value=begin_cm)
+
+                def default_query_side_effect(*args):
+                    query = MagicMock()
+                    if args and args[0] == Dataset and shared_mock_data["dataset"] is not None:
+                        where_result = MagicMock()
+                        where_result.first.return_value = shared_mock_data["dataset"]
+                        query.where = MagicMock(return_value=where_result)
+                    elif args and args[0] == Document and shared_mock_data["documents"] is not None:
+                        where_result = MagicMock()
+                        where_result.where = MagicMock(return_value=where_result)
+
+                        if shared_mock_data["doc_iter"] is None:
+                            docs = shared_mock_data["documents"] or [None]
+                            shared_mock_data["doc_iter"] = iter(docs)
+
+                        where_result.first = lambda: next(shared_mock_data["doc_iter"], None)
+                        docs_or_empty = shared_mock_data["documents"] or []
+                        where_result.all = MagicMock(return_value=docs_or_empty)
+                        query.where = MagicMock(return_value=where_result)
+                    else:
+                        query.where = MagicMock(return_value=query)
+                    return query
+
+                self._default_session.query = MagicMock(side_effect=default_query_side_effect)
+
+            def __getattr__(self, name):
+                # Forward all attribute access to the first session, or default if none created yet
+                target_session = self._sessions[0] if self._sessions else self._default_session
+                return getattr(target_session, name)
+
+            @property
+            def all_sessions(self):
+                """Access all created sessions for testing."""
+                return self._sessions
+
+        wrapper = SessionWrapper()
+        yield wrapper
 
 
 @pytest.fixture
@@ -252,18 +356,9 @@ class TestTaskEnqueuing:
         use the deprecated function.
         """
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                # Return documents one by one for each call
-                mock_query.where.return_value.first.side_effect = mock_documents
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -304,21 +399,9 @@ class TestBatchProcessing:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        # Create an iterator for documents
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                # Return documents one by one for each call
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -357,19 +440,9 @@ class TestBatchProcessing:
             doc.stopped_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         mock_feature_service.get_features.return_value.billing.enabled = True
         mock_feature_service.get_features.return_value.billing.subscription.plan = CloudPlan.PROFESSIONAL
@@ -407,19 +480,9 @@ class TestBatchProcessing:
             doc.stopped_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         mock_feature_service.get_features.return_value.billing.enabled = True
         mock_feature_service.get_features.return_value.billing.subscription.plan = CloudPlan.SANDBOX
@@ -444,7 +507,10 @@ class TestBatchProcessing:
         """
         # Arrange
         document_ids = []
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
+
+        # Set shared mock data with empty documents list
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = []
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -482,19 +548,9 @@ class TestProgressTracking:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -528,19 +584,9 @@ class TestProgressTracking:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -635,19 +681,9 @@ class TestErrorHandling:
             doc.stopped_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Set up to trigger vector space limit error
         mock_feature_service.get_features.return_value.billing.enabled = True
@@ -674,17 +710,9 @@ class TestErrorHandling:
         Errors during indexing should be caught and logged, but not crash the task.
         """
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first.side_effect = mock_documents
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Make IndexingRunner raise an exception
         mock_indexing_runner.run.side_effect = Exception("Indexing failed")
@@ -708,17 +736,9 @@ class TestErrorHandling:
         but not treated as a failure.
         """
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first.side_effect = mock_documents
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Make IndexingRunner raise DocumentIsPausedError
         mock_indexing_runner.run.side_effect = DocumentIsPausedError("Document is paused")
@@ -853,17 +873,9 @@ class TestTaskCancellation:
         Session cleanup should happen in finally block.
         """
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first.side_effect = mock_documents
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -883,17 +895,9 @@ class TestTaskCancellation:
         Session cleanup should happen even when errors occur.
         """
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first.side_effect = mock_documents
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Make IndexingRunner raise an exception
         mock_indexing_runner.run.side_effect = Exception("Test error")
@@ -962,6 +966,7 @@ class TestAdvancedScenarios:
         document_ids = [str(uuid.uuid4()) for _ in range(3)]
 
         # Create only 2 documents (simulate one missing)
+        # The new code uses .all() which will only return existing documents
         mock_documents = []
         for i, doc_id in enumerate([document_ids[0], document_ids[2]]):  # Skip middle one
             doc = MagicMock(spec=Document)
@@ -971,21 +976,9 @@ class TestAdvancedScenarios:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        # Create iterator that returns None for missing document
-        doc_responses = [mock_documents[0], None, mock_documents[1]]
-        doc_iter = iter(doc_responses)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data - .all() will only return existing documents
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -1075,19 +1068,9 @@ class TestAdvancedScenarios:
             doc.stopped_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Set vector space exactly at limit
         mock_feature_service.get_features.return_value.billing.enabled = True
@@ -1219,19 +1202,9 @@ class TestAdvancedScenarios:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Billing disabled - limits should not be checked
         mock_feature_service.get_features.return_value.billing.enabled = False
@@ -1273,19 +1246,9 @@ class TestIntegration:
 
         # Set up rpop to return None for concurrency check (no more tasks)
         mock_redis.rpop.side_effect = [None]
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -1321,19 +1284,9 @@ class TestIntegration:
 
         # Set up rpop to return None for concurrency check (no more tasks)
         mock_redis.rpop.side_effect = [None]
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -1415,17 +1368,9 @@ class TestEdgeCases:
         mock_document.indexing_status = "waiting"
         mock_document.processing_started_at = None
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: mock_document
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = [mock_document]
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -1465,17 +1410,9 @@ class TestEdgeCases:
         mock_document.indexing_status = "waiting"
         mock_document.processing_started_at = None
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: mock_document
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = [mock_document]
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -1555,19 +1492,9 @@ class TestEdgeCases:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Set vector space limit to 0 (unlimited)
         mock_feature_service.get_features.return_value.billing.enabled = True
@@ -1612,19 +1539,9 @@ class TestEdgeCases:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Set negative vector space limit
         mock_feature_service.get_features.return_value.billing.enabled = True
@@ -1675,19 +1592,9 @@ class TestPerformanceScenarios:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Configure billing with sufficient limits
         mock_feature_service.get_features.return_value.billing.enabled = True
@@ -1826,19 +1733,9 @@ class TestRobustness:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         # Make IndexingRunner raise an exception
         mock_indexing_runner.run.side_effect = RuntimeError("Unexpected indexing error")
@@ -1866,7 +1763,7 @@ class TestRobustness:
         - No exceptions occur
 
         Expected behavior:
-        - Database session is closed
+        - All database sessions are closed
         - No connection leaks
         """
         # Arrange
@@ -1879,19 +1776,9 @@ class TestRobustness:
             doc.processing_started_at = None
             mock_documents.append(doc)
 
-        mock_db_session.query.return_value.where.return_value.first.return_value = mock_dataset
-
-        doc_iter = iter(mock_documents)
-
-        def mock_query_side_effect(*args):
-            mock_query = MagicMock()
-            if args[0] == Dataset:
-                mock_query.where.return_value.first.return_value = mock_dataset
-            elif args[0] == Document:
-                mock_query.where.return_value.first = lambda: next(doc_iter, None)
-            return mock_query
-
-        mock_db_session.query.side_effect = mock_query_side_effect
+        # Set shared mock data so all sessions can access it
+        mock_db_session._shared_data["dataset"] = mock_dataset
+        mock_db_session._shared_data["documents"] = mock_documents
 
         with patch("tasks.document_indexing_task.FeatureService.get_features") as mock_features:
             mock_features.return_value.billing.enabled = False
@@ -1899,10 +1786,11 @@ class TestRobustness:
             # Act
             _document_indexing(dataset_id, document_ids)
 
-            # Assert
-            assert mock_db_session.close.called
-            # Verify close is called exactly once
-            assert mock_db_session.close.call_count == 1
+            # Assert - All created sessions should be closed
+            # The code creates multiple sessions: validation, Phase 1 (parsing), Phase 3 (summary)
+            assert len(mock_db_session.all_sessions) >= 1
+            for session in mock_db_session.all_sessions:
+                assert session.close.called, "All sessions should be closed"
 
     def test_task_proxy_handles_feature_service_failure(self, tenant_id, dataset_id, document_ids, mock_redis):
         """

api/tests/unit_tests/tasks/test_document_indexing_sync_task.py

@@ -109,25 +109,87 @@ def mock_document_segments(document_id):
 
 @pytest.fixture
 def mock_db_session():
-    """Mock database session via session_factory.create_session()."""
+    """Mock database session via session_factory.create_session().
+
+    After session split refactor, the code calls create_session() multiple times.
+    This fixture creates shared query mocks so all sessions use the same
+    query configuration, simulating database persistence across sessions.
+
+    The fixture automatically converts side_effect to cycle to prevent StopIteration.
+    Tests configure mocks the same way as before, but behind the scenes the values
+    are cycled infinitely for all sessions.
+    """
+    from itertools import cycle
+
     with patch("tasks.document_indexing_sync_task.session_factory") as mock_sf:
-        session = MagicMock()
-        # Ensure tests can observe session.close() via context manager teardown
-        session.close = MagicMock()
-        cm = MagicMock()
-        cm.__enter__.return_value = session
+        sessions = []
 
-        def _exit_side_effect(*args, **kwargs):
-            session.close()
+        # Shared query mocks - all sessions use these
+        shared_query = MagicMock()
+        shared_filter_by = MagicMock()
+        shared_scalars_result = MagicMock()
 
-        cm.__exit__.side_effect = _exit_side_effect
-        mock_sf.create_session.return_value = cm
+        # Create custom first mock that auto-cycles side_effect
+        class CyclicMock(MagicMock):
+            def __setattr__(self, name, value):
+                if name == "side_effect" and value is not None:
+                    # Convert list/tuple to infinite cycle
+                    if isinstance(value, (list, tuple)):
+                        value = cycle(value)
+                super().__setattr__(name, value)
 
-        query = MagicMock()
-        session.query.return_value = query
-        query.where.return_value = query
-        session.scalars.return_value = MagicMock()
-        yield session
+        shared_query.where.return_value.first = CyclicMock()
+        shared_filter_by.first = CyclicMock()
+
+        def _create_session():
+            """Create a new mock session for each create_session() call."""
+            session = MagicMock()
+            session.close = MagicMock()
+            session.commit = MagicMock()
+
+            # Mock session.begin() context manager
+            begin_cm = MagicMock()
+            begin_cm.__enter__.return_value = session
+
+            def _begin_exit_side_effect(exc_type, exc, tb):
+                # commit on success
+                if exc_type is None:
+                    session.commit()
+                # return False to propagate exceptions
+                return False
+
+            begin_cm.__exit__.side_effect = _begin_exit_side_effect
+            session.begin.return_value = begin_cm
+
+            # Mock create_session() context manager
+            cm = MagicMock()
+            cm.__enter__.return_value = session
+
+            def _exit_side_effect(exc_type, exc, tb):
+                session.close()
+                return False
+
+            cm.__exit__.side_effect = _exit_side_effect
+
+            # All sessions use the same shared query mocks
+            session.query.return_value = shared_query
+            shared_query.where.return_value = shared_query
+            shared_query.filter_by.return_value = shared_filter_by
+            session.scalars.return_value = shared_scalars_result
+
+            sessions.append(session)
+            # Attach helpers on the first created session for assertions across all sessions
+            if len(sessions) == 1:
+                session.get_all_sessions = lambda: sessions
+                session.any_close_called = lambda: any(s.close.called for s in sessions)
+                session.any_commit_called = lambda: any(s.commit.called for s in sessions)
+            return cm
+
+        mock_sf.create_session.side_effect = _create_session
+
+        # Create first session and return it
+        _create_session()
+        yield sessions[0]
 
 
 @pytest.fixture
@@ -186,8 +248,8 @@ class TestDocumentIndexingSyncTask:
         # Act
         document_indexing_sync_task(dataset_id, document_id)
 
-        # Assert
-        mock_db_session.close.assert_called_once()
+        # Assert - at least one session should have been closed
+        assert mock_db_session.any_close_called()
 
     def test_missing_notion_workspace_id(self, mock_db_session, mock_document, dataset_id, document_id):
         """Test that task raises error when notion_workspace_id is missing."""
@@ -230,6 +292,7 @@ class TestDocumentIndexingSyncTask:
         """Test that task handles missing credentials by updating document status."""
         # Arrange
         mock_db_session.query.return_value.where.return_value.first.return_value = mock_document
+        mock_db_session.query.return_value.filter_by.return_value.first.return_value = mock_document
         mock_datasource_provider_service.get_datasource_credentials.return_value = None
 
         # Act
@@ -239,8 +302,8 @@ class TestDocumentIndexingSyncTask:
         assert mock_document.indexing_status == "error"
         assert "Datasource credential not found" in mock_document.error
         assert mock_document.stopped_at is not None
-        mock_db_session.commit.assert_called()
-        mock_db_session.close.assert_called()
+        assert mock_db_session.any_commit_called()
+        assert mock_db_session.any_close_called()
 
     def test_page_not_updated(
         self,
@@ -254,6 +317,7 @@ class TestDocumentIndexingSyncTask:
         """Test that task does nothing when page has not been updated."""
         # Arrange
         mock_db_session.query.return_value.where.return_value.first.return_value = mock_document
+        mock_db_session.query.return_value.filter_by.return_value.first.return_value = mock_document
         # Return same time as stored in document
         mock_notion_extractor.get_notion_last_edited_time.return_value = "2024-01-01T00:00:00Z"
 
@@ -263,8 +327,8 @@ class TestDocumentIndexingSyncTask:
         # Assert
         # Document status should remain unchanged
         assert mock_document.indexing_status == "completed"
-        # Session should still be closed via context manager teardown
-        assert mock_db_session.close.called
+        # At least one session should have been closed via context manager teardown
+        assert mock_db_session.any_close_called()
 
     def test_successful_sync_when_page_updated(
         self,
@@ -281,7 +345,20 @@ class TestDocumentIndexingSyncTask:
     ):
         """Test successful sync flow when Notion page has been updated."""
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.side_effect = [mock_document, mock_dataset]
+        # Set exact sequence of returns across calls to `.first()`:
+        # 1) document (initial fetch)
+        # 2) dataset (pre-check)
+        # 3) dataset (cleaning phase)
+        # 4) document (pre-indexing update)
+        # 5) document (indexing runner fetch)
+        mock_db_session.query.return_value.where.return_value.first.side_effect = [
+            mock_document,
+            mock_dataset,
+            mock_dataset,
+            mock_document,
+            mock_document,
+        ]
+        mock_db_session.query.return_value.filter_by.return_value.first.return_value = mock_document
         mock_db_session.scalars.return_value.all.return_value = mock_document_segments
         # NotionExtractor returns updated time
         mock_notion_extractor.get_notion_last_edited_time.return_value = "2024-01-02T00:00:00Z"
@@ -299,28 +376,40 @@ class TestDocumentIndexingSyncTask:
         mock_processor.clean.assert_called_once()
 
         # Verify segments were deleted from database in batch (DELETE FROM document_segments)
-        execute_sqls = [" ".join(str(c[0][0]).split()) for c in mock_db_session.execute.call_args_list]
+        # Aggregate execute calls across all created sessions
+        execute_sqls = []
+        for s in mock_db_session.get_all_sessions():
+            execute_sqls.extend([" ".join(str(c[0][0]).split()) for c in s.execute.call_args_list])
         assert any("DELETE FROM document_segments" in sql for sql in execute_sqls)
 
         # Verify indexing runner was called
         mock_indexing_runner.run.assert_called_once_with([mock_document])
 
-        # Verify session operations
-        assert mock_db_session.commit.called
-        mock_db_session.close.assert_called_once()
+        # Verify session operations (across any created session)
+        assert mock_db_session.any_commit_called()
+        assert mock_db_session.any_close_called()
 
     def test_dataset_not_found_during_cleaning(
         self,
         mock_db_session,
         mock_datasource_provider_service,
         mock_notion_extractor,
+        mock_indexing_runner,
         mock_document,
         dataset_id,
         document_id,
     ):
         """Test that task handles dataset not found during cleaning phase."""
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.side_effect = [mock_document, None]
+        # Sequence: document (initial), dataset (pre-check), None (cleaning), document (update), document (indexing)
+        mock_db_session.query.return_value.where.return_value.first.side_effect = [
+            mock_document,
+            mock_dataset,
+            None,
+            mock_document,
+            mock_document,
+        ]
+        mock_db_session.query.return_value.filter_by.return_value.first.return_value = mock_document
         mock_notion_extractor.get_notion_last_edited_time.return_value = "2024-01-02T00:00:00Z"
 
         # Act
@@ -329,8 +418,8 @@ class TestDocumentIndexingSyncTask:
         # Assert
         # Document should still be set to parsing
         assert mock_document.indexing_status == "parsing"
-        # Session should be closed after error
-        mock_db_session.close.assert_called_once()
+        # At least one session should be closed after error
+        assert mock_db_session.any_close_called()
 
     def test_cleaning_error_continues_to_indexing(
         self,
@@ -346,8 +435,14 @@ class TestDocumentIndexingSyncTask:
     ):
         """Test that indexing continues even if cleaning fails."""
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.side_effect = [mock_document, mock_dataset]
-        mock_db_session.scalars.return_value.all.side_effect = Exception("Cleaning error")
+        from itertools import cycle
+
+        mock_db_session.query.return_value.where.return_value.first.side_effect = cycle([mock_document, mock_dataset])
+        mock_db_session.query.return_value.filter_by.return_value.first.return_value = mock_document
+        # Make the cleaning step fail but not the segment fetch
+        processor = mock_index_processor_factory.return_value.init_index_processor.return_value
+        processor.clean.side_effect = Exception("Cleaning error")
+        mock_db_session.scalars.return_value.all.return_value = []
         mock_notion_extractor.get_notion_last_edited_time.return_value = "2024-01-02T00:00:00Z"
 
         # Act
@@ -356,7 +451,7 @@ class TestDocumentIndexingSyncTask:
         # Assert
         # Indexing should still be attempted despite cleaning error
         mock_indexing_runner.run.assert_called_once_with([mock_document])
-        mock_db_session.close.assert_called_once()
+        assert mock_db_session.any_close_called()
 
     def test_indexing_runner_document_paused_error(
         self,
@@ -373,7 +468,10 @@ class TestDocumentIndexingSyncTask:
     ):
         """Test that DocumentIsPausedError is handled gracefully."""
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.side_effect = [mock_document, mock_dataset]
+        from itertools import cycle
+
+        mock_db_session.query.return_value.where.return_value.first.side_effect = cycle([mock_document, mock_dataset])
+        mock_db_session.query.return_value.filter_by.return_value.first.return_value = mock_document
         mock_db_session.scalars.return_value.all.return_value = mock_document_segments
         mock_notion_extractor.get_notion_last_edited_time.return_value = "2024-01-02T00:00:00Z"
         mock_indexing_runner.run.side_effect = DocumentIsPausedError("Document paused")
@@ -383,7 +481,7 @@ class TestDocumentIndexingSyncTask:
 
         # Assert
         # Session should be closed after handling error
-        mock_db_session.close.assert_called_once()
+        assert mock_db_session.any_close_called()
 
     def test_indexing_runner_general_error(
         self,
@@ -400,7 +498,10 @@ class TestDocumentIndexingSyncTask:
     ):
         """Test that general exceptions during indexing are handled."""
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.side_effect = [mock_document, mock_dataset]
+        from itertools import cycle
+
+        mock_db_session.query.return_value.where.return_value.first.side_effect = cycle([mock_document, mock_dataset])
+        mock_db_session.query.return_value.filter_by.return_value.first.return_value = mock_document
         mock_db_session.scalars.return_value.all.return_value = mock_document_segments
         mock_notion_extractor.get_notion_last_edited_time.return_value = "2024-01-02T00:00:00Z"
         mock_indexing_runner.run.side_effect = Exception("Indexing error")
@@ -410,7 +511,7 @@ class TestDocumentIndexingSyncTask:
 
         # Assert
         # Session should be closed after error
-        mock_db_session.close.assert_called_once()
+        assert mock_db_session.any_close_called()
 
     def test_notion_extractor_initialized_with_correct_params(
         self,
@@ -517,7 +618,14 @@ class TestDocumentIndexingSyncTask:
     ):
         """Test that index processor clean is called with correct parameters."""
         # Arrange
-        mock_db_session.query.return_value.where.return_value.first.side_effect = [mock_document, mock_dataset]
+        # Sequence: document (initial), dataset (pre-check), dataset (cleaning), document (update), document (indexing)
+        mock_db_session.query.return_value.where.return_value.first.side_effect = [
+            mock_document,
+            mock_dataset,
+            mock_dataset,
+            mock_document,
+            mock_document,
+        ]
         mock_db_session.scalars.return_value.all.return_value = mock_document_segments
         mock_notion_extractor.get_notion_last_edited_time.return_value = "2024-01-02T00:00:00Z"
 

api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py

@@ -350,7 +350,7 @@ class TestDeleteWorkflowArchiveLogs:
         mock_query.where.return_value = mock_delete_query
         mock_db.session.query.return_value = mock_query
 
-        delete_func("log-1")
+        delete_func(mock_db.session, "log-1")
 
         mock_db.session.query.assert_called_once_with(WorkflowArchiveLog)
         mock_query.where.assert_called_once()