test: migrate remaining DocumentSegment navigation SQL tests to testcontainers (#32523)
Some checks are pending
autofix.ci / autofix (push) Waiting to run
Build and Push API & Web / build (api, DIFY_API_IMAGE_NAME, linux/amd64, build-api-amd64) (push) Waiting to run
Build and Push API & Web / build (api, DIFY_API_IMAGE_NAME, linux/arm64, build-api-arm64) (push) Waiting to run
Build and Push API & Web / build (web, DIFY_WEB_IMAGE_NAME, linux/amd64, build-web-amd64) (push) Waiting to run
Build and Push API & Web / build (web, DIFY_WEB_IMAGE_NAME, linux/arm64, build-web-arm64) (push) Waiting to run
Build and Push API & Web / create-manifest (api, DIFY_API_IMAGE_NAME, merge-api-images) (push) Blocked by required conditions
Build and Push API & Web / create-manifest (web, DIFY_WEB_IMAGE_NAME, merge-web-images) (push) Blocked by required conditions
Main CI Pipeline / Check Changed Files (push) Waiting to run
Main CI Pipeline / API Tests (push) Blocked by required conditions
Main CI Pipeline / Web Tests (push) Blocked by required conditions
Main CI Pipeline / Style Check (push) Waiting to run
Main CI Pipeline / VDB Tests (push) Blocked by required conditions
Main CI Pipeline / DB Migration Test (push) Blocked by required conditions

Co-authored-by: KinomotoMio <200703522+KinomotoMio@users.noreply.github.com>
This commit is contained in:
木之本澪
2026-02-25 01:51:38 +08:00
committed by GitHub
parent beea1acd92
commit 3abfbc0246
2 changed files with 218 additions and 142 deletions

View File

@@ -269,3 +269,221 @@ class TestDatasetDocumentProperties:
db_session_with_containers.flush()
assert doc.hit_count == 25
class TestDocumentSegmentNavigationProperties:
"""Integration tests for DocumentSegment navigation properties."""
@pytest.fixture(autouse=True)
def _auto_rollback(self, db_session_with_containers: Session) -> Generator[None, None, None]:
"""Automatically rollback session changes after each test."""
yield
db_session_with_containers.rollback()
def test_document_segment_dataset_property(self, db_session_with_containers: Session) -> None:
"""Test segment can access its parent dataset."""
# Arrange
tenant_id = str(uuid4())
created_by = str(uuid4())
dataset = Dataset(
tenant_id=tenant_id,
name="Test Dataset",
data_source_type="upload_file",
created_by=created_by,
)
db_session_with_containers.add(dataset)
db_session_with_containers.flush()
document = Document(
tenant_id=tenant_id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch="batch_001",
name="test.pdf",
created_from="web",
created_by=created_by,
)
db_session_with_containers.add(document)
db_session_with_containers.flush()
segment = DocumentSegment(
tenant_id=tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content="Test",
word_count=1,
tokens=2,
created_by=created_by,
)
db_session_with_containers.add(segment)
db_session_with_containers.flush()
# Act
related_dataset = segment.dataset
# Assert
assert related_dataset is not None
assert related_dataset.id == dataset.id
def test_document_segment_document_property(self, db_session_with_containers: Session) -> None:
"""Test segment can access its parent document."""
# Arrange
tenant_id = str(uuid4())
created_by = str(uuid4())
dataset = Dataset(
tenant_id=tenant_id,
name="Test Dataset",
data_source_type="upload_file",
created_by=created_by,
)
db_session_with_containers.add(dataset)
db_session_with_containers.flush()
document = Document(
tenant_id=tenant_id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch="batch_001",
name="test.pdf",
created_from="web",
created_by=created_by,
)
db_session_with_containers.add(document)
db_session_with_containers.flush()
segment = DocumentSegment(
tenant_id=tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content="Test",
word_count=1,
tokens=2,
created_by=created_by,
)
db_session_with_containers.add(segment)
db_session_with_containers.flush()
# Act
related_document = segment.document
# Assert
assert related_document is not None
assert related_document.id == document.id
def test_document_segment_previous_segment(self, db_session_with_containers: Session) -> None:
"""Test segment can access previous segment."""
# Arrange
tenant_id = str(uuid4())
created_by = str(uuid4())
dataset = Dataset(
tenant_id=tenant_id,
name="Test Dataset",
data_source_type="upload_file",
created_by=created_by,
)
db_session_with_containers.add(dataset)
db_session_with_containers.flush()
document = Document(
tenant_id=tenant_id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch="batch_001",
name="test.pdf",
created_from="web",
created_by=created_by,
)
db_session_with_containers.add(document)
db_session_with_containers.flush()
previous_segment = DocumentSegment(
tenant_id=tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content="Previous",
word_count=1,
tokens=2,
created_by=created_by,
)
segment = DocumentSegment(
tenant_id=tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=2,
content="Current",
word_count=1,
tokens=2,
created_by=created_by,
)
db_session_with_containers.add_all([previous_segment, segment])
db_session_with_containers.flush()
# Act
prev_seg = segment.previous_segment
# Assert
assert prev_seg is not None
assert prev_seg.position == 1
def test_document_segment_next_segment(self, db_session_with_containers: Session) -> None:
"""Test segment can access next segment."""
# Arrange
tenant_id = str(uuid4())
created_by = str(uuid4())
dataset = Dataset(
tenant_id=tenant_id,
name="Test Dataset",
data_source_type="upload_file",
created_by=created_by,
)
db_session_with_containers.add(dataset)
db_session_with_containers.flush()
document = Document(
tenant_id=tenant_id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch="batch_001",
name="test.pdf",
created_from="web",
created_by=created_by,
)
db_session_with_containers.add(document)
db_session_with_containers.flush()
segment = DocumentSegment(
tenant_id=tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content="Current",
word_count=1,
tokens=2,
created_by=created_by,
)
next_segment = DocumentSegment(
tenant_id=tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=2,
content="Next",
word_count=1,
tokens=2,
created_by=created_by,
)
db_session_with_containers.add_all([segment, next_segment])
db_session_with_containers.flush()
# Act
next_seg = segment.next_segment
# Assert
assert next_seg is not None
assert next_seg.position == 2

View File

@@ -954,148 +954,6 @@ class TestChildChunk:
assert child_chunk.index_node_hash == index_node_hash
class TestDocumentSegmentNavigation:
"""Test suite for DocumentSegment navigation properties."""
def test_document_segment_dataset_property(self):
"""Test segment can access its parent dataset."""
# Arrange
dataset_id = str(uuid4())
segment = DocumentSegment(
tenant_id=str(uuid4()),
dataset_id=dataset_id,
document_id=str(uuid4()),
position=1,
content="Test",
word_count=1,
tokens=2,
created_by=str(uuid4()),
)
mock_dataset = Dataset(
tenant_id=str(uuid4()),
name="Test Dataset",
data_source_type="upload_file",
created_by=str(uuid4()),
)
mock_dataset.id = dataset_id
# Mock the database session scalar
with patch("models.dataset.db.session.scalar", return_value=mock_dataset):
# Act
dataset = segment.dataset
# Assert
assert dataset is not None
assert dataset.id == dataset_id
def test_document_segment_document_property(self):
"""Test segment can access its parent document."""
# Arrange
document_id = str(uuid4())
segment = DocumentSegment(
tenant_id=str(uuid4()),
dataset_id=str(uuid4()),
document_id=document_id,
position=1,
content="Test",
word_count=1,
tokens=2,
created_by=str(uuid4()),
)
mock_document = Document(
tenant_id=str(uuid4()),
dataset_id=str(uuid4()),
position=1,
data_source_type="upload_file",
batch="batch_001",
name="test.pdf",
created_from="web",
created_by=str(uuid4()),
)
mock_document.id = document_id
# Mock the database session scalar
with patch("models.dataset.db.session.scalar", return_value=mock_document):
# Act
document = segment.document
# Assert
assert document is not None
assert document.id == document_id
def test_document_segment_previous_segment(self):
"""Test segment can access previous segment."""
# Arrange
document_id = str(uuid4())
segment = DocumentSegment(
tenant_id=str(uuid4()),
dataset_id=str(uuid4()),
document_id=document_id,
position=2,
content="Test",
word_count=1,
tokens=2,
created_by=str(uuid4()),
)
previous_segment = DocumentSegment(
tenant_id=str(uuid4()),
dataset_id=str(uuid4()),
document_id=document_id,
position=1,
content="Previous",
word_count=1,
tokens=2,
created_by=str(uuid4()),
)
# Mock the database session scalar
with patch("models.dataset.db.session.scalar", return_value=previous_segment):
# Act
prev_seg = segment.previous_segment
# Assert
assert prev_seg is not None
assert prev_seg.position == 1
def test_document_segment_next_segment(self):
"""Test segment can access next segment."""
# Arrange
document_id = str(uuid4())
segment = DocumentSegment(
tenant_id=str(uuid4()),
dataset_id=str(uuid4()),
document_id=document_id,
position=1,
content="Test",
word_count=1,
tokens=2,
created_by=str(uuid4()),
)
next_segment = DocumentSegment(
tenant_id=str(uuid4()),
dataset_id=str(uuid4()),
document_id=document_id,
position=2,
content="Next",
word_count=1,
tokens=2,
created_by=str(uuid4()),
)
# Mock the database session scalar
with patch("models.dataset.db.session.scalar", return_value=next_segment):
# Act
next_seg = segment.next_segment
# Assert
assert next_seg is not None
assert next_seg.position == 2
class TestModelIntegration:
"""Test suite for model integration scenarios."""