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>
This commit is contained in:
bohdansolovie
2026-04-14 03:56:11 -04:00
committed by CodingOnStar
parent 6c382ddb05
commit 92c4725f4b
2 changed files with 10 additions and 8 deletions

View File

@ -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

View File

@ -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