From a951cc996b234660c2498d2c0a093a589ab0c1ff Mon Sep 17 00:00:00 2001
From: bohdansolovie <153934212+bohdansolovie@users.noreply.github.com>
Date: Tue, 14 Apr 2026 03:56:11 -0400
Subject: [PATCH] test: migrate document indexing task tests to SQLAlchemy 2.0
 select API (#35145)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 .../tasks/test_disable_segments_from_index_task.py    |  7 ++++---
 .../tasks/test_document_indexing_update_task.py       | 11 ++++++-----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py
index 3e9a0c8f7f..6e03bd9351 100644
--- a/api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py
+++ b/api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py
@@ -9,6 +9,7 @@ The task is responsible for removing document segments from the search index whe
 from unittest.mock import MagicMock, patch
 
 from faker import Faker
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
@@ -471,9 +472,9 @@ class TestDisableSegmentsFromIndexTask:
                 db_session_with_containers.refresh(segments[1])
 
                 # Check that segments are re-enabled after error
-                updated_segments = (
-                    db_session_with_containers.query(DocumentSegment).where(DocumentSegment.id.in_(segment_ids)).all()
-                )
+                updated_segments = db_session_with_containers.scalars(
+                    select(DocumentSegment).where(DocumentSegment.id.in_(segment_ids))
+                ).all()
 
                 for segment in updated_segments:
                     assert segment.enabled is True
diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py
index d94abf2b40..a9a8c0f30c 100644
--- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py
+++ b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py
@@ -2,6 +2,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 from faker import Faker
+from sqlalchemy import func, select
 
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
@@ -123,13 +124,13 @@ class TestDocumentIndexingUpdateTask:
         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()
+        updated = db_session_with_containers.scalar(select(Document).where(Document.id == document.id).limit(1))
         assert updated.indexing_status == IndexingStatus.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()
+        remaining = db_session_with_containers.scalar(
+            select(func.count()).select_from(DocumentSegment).where(DocumentSegment.document_id == document.id)
         )
         assert remaining == 0
 
@@ -167,8 +168,8 @@ class TestDocumentIndexingUpdateTask:
         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()
+        remaining = db_session_with_containers.scalar(
+            select(func.count()).select_from(DocumentSegment).where(DocumentSegment.document_id == document.id)
         )
         assert remaining > 0