"""
Integration tests for clean_dataset_task using testcontainers.
This module provides comprehensive integration tests for the dataset cleanup task
using TestContainers infrastructure. The tests ensure that the task properly
cleans up all dataset-related data including vector indexes, documents,
segments, metadata, and storage files in a real database environment.
All tests use the testcontainers infrastructure to ensure proper database isolation
and realistic testing scenarios with actual PostgreSQL and Redis instances.
"""
import uuid
from datetime import datetime
from unittest.mock import MagicMock, patch
import pytest
from faker import Faker
from sqlalchemy import delete, select
from sqlalchemy.orm import Session
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from extensions.storage.storage_type import StorageType
from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import (
AppDatasetJoin,
Dataset,
DatasetMetadata,
DatasetMetadataBinding,
DatasetProcessRule,
DatasetQuery,
Document,
DocumentSegment,
)
from models.enums import (
CreatorUserRole,
DatasetMetadataType,
DataSourceType,
DocumentCreatedFrom,
IndexingStatus,
SegmentStatus,
)
from models.model import UploadFile
from tasks.clean_dataset_task import clean_dataset_task
class TestCleanDatasetTask:
"""Integration tests for clean_dataset_task using testcontainers."""
@pytest.fixture(autouse=True)
def cleanup_database(self, db_session_with_containers: Session):
"""Clean up database before each test to ensure isolation."""
from extensions.ext_redis import redis_client
# Clear all test data using the provided session fixture
db_session_with_containers.execute(delete(DatasetMetadataBinding))
db_session_with_containers.execute(delete(DatasetMetadata))
db_session_with_containers.execute(delete(AppDatasetJoin))
db_session_with_containers.execute(delete(DatasetQuery))
db_session_with_containers.execute(delete(DatasetProcessRule))
db_session_with_containers.execute(delete(DocumentSegment))
db_session_with_containers.execute(delete(Document))
db_session_with_containers.execute(delete(Dataset))
db_session_with_containers.execute(delete(UploadFile))
db_session_with_containers.execute(delete(TenantAccountJoin))
db_session_with_containers.execute(delete(Tenant))
db_session_with_containers.execute(delete(Account))
db_session_with_containers.commit()
# Clear Redis cache
redis_client.flushdb()
@pytest.fixture
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.clean_dataset_task.storage", autospec=True) as mock_storage,
patch("tasks.clean_dataset_task.IndexProcessorFactory", autospec=True) as mock_index_processor_factory,
):
# Setup default mock returns
mock_storage.delete.return_value = None
# Mock index processor
mock_index_processor = MagicMock()
mock_index_processor.clean.return_value = None
mock_index_processor_factory_instance = MagicMock()
mock_index_processor_factory_instance.init_index_processor.return_value = mock_index_processor
mock_index_processor_factory.return_value = mock_index_processor_factory_instance
yield {
"storage": mock_storage,
"index_processor_factory": mock_index_processor_factory,
"index_processor": mock_index_processor,
}
def _create_test_account_and_tenant(self, db_session_with_containers: Session):
"""
Helper method to create a test account and tenant for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
Returns:
tuple: (Account, Tenant) created instances
"""
fake = Faker()
# Create account
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()
# Create tenant
tenant = Tenant(
name=fake.company(),
plan="basic",
status="normal",
)
db_session_with_containers.add(tenant)
db_session_with_containers.commit()
# Create tenant-account relationship
tenant_account_join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER,
)
db_session_with_containers.add(tenant_account_join)
db_session_with_containers.commit()
return account, tenant
def _create_test_dataset(self, db_session_with_containers: Session, account, tenant):
"""
Helper method to create a test dataset for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
tenant: Tenant instance
Returns:
Dataset: Created dataset instance
"""
dataset = Dataset(
id=str(uuid.uuid4()),
tenant_id=tenant.id,
name="test_dataset",
description="Test dataset for cleanup testing",
indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph"}',
collection_binding_id=str(uuid.uuid4()),
created_by=account.id,
created_at=datetime.now(),
updated_at=datetime.now(),
)
db_session_with_containers.add(dataset)
db_session_with_containers.commit()
return dataset
def _create_test_document(self, db_session_with_containers: Session, account, tenant, dataset):
"""
Helper method to create a test document for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
tenant: Tenant instance
dataset: Dataset instance
Returns:
Document: Created document instance
"""
document = Document(
id=str(uuid.uuid4()),
tenant_id=tenant.id,
dataset_id=dataset.id,
position=1,
data_source_type=DataSourceType.UPLOAD_FILE,
batch="test_batch",
name="test_document",
created_from=DocumentCreatedFrom.WEB,
created_by=account.id,
indexing_status=IndexingStatus.COMPLETED,
enabled=True,
archived=False,
doc_form=IndexStructureType.PARAGRAPH_INDEX,
word_count=100,
created_at=datetime.now(),
updated_at=datetime.now(),
)
db_session_with_containers.add(document)
db_session_with_containers.commit()
return document
def _create_test_segment(self, db_session_with_containers: Session, account, tenant, dataset, document):
"""
Helper method to create a test document segment for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
tenant: Tenant instance
dataset: Dataset instance
document: Document instance
Returns:
DocumentSegment: Created document segment instance
"""
segment = DocumentSegment(
tenant_id=tenant.id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content="This is a test segment content for cleanup testing",
word_count=20,
tokens=30,
created_by=account.id,
status=SegmentStatus.COMPLETED,
index_node_id=str(uuid.uuid4()),
index_node_hash="test_hash",
)
db_session_with_containers.add(segment)
db_session_with_containers.commit()
return segment
def _create_test_upload_file(self, db_session_with_containers: Session, account, tenant):
"""
Helper method to create a test upload file for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
tenant: Tenant instance
Returns:
UploadFile: Created upload file instance
"""
fake = Faker()
upload_file = UploadFile(
tenant_id=tenant.id,
storage_type=StorageType.LOCAL,
key=f"test_files/{fake.file_name()}",
name=fake.file_name(),
size=1024,
extension=".txt",
mime_type="text/plain",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(),
used=False,
)
db_session_with_containers.add(upload_file)
db_session_with_containers.commit()
return upload_file
def test_clean_dataset_task_success_basic_cleanup(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test successful basic dataset cleanup with minimal data.
This test verifies that the task can successfully:
1. Clean up vector database indexes
2. Delete documents and segments
3. Remove dataset metadata and bindings
4. Handle empty document scenarios
5. Complete cleanup process without errors
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
# Execute the task
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
)
# Verify results
# Check that dataset-related data was cleaned up
documents = db_session_with_containers.scalars(select(Document).where(Document.dataset_id == dataset.id)).all()
assert len(documents) == 0
segments = db_session_with_containers.scalars(
select(DocumentSegment).where(DocumentSegment.dataset_id == dataset.id)
).all()
assert len(segments) == 0
# Check that metadata and bindings were cleaned up
metadata = db_session_with_containers.scalars(
select(DatasetMetadata).where(DatasetMetadata.dataset_id == dataset.id)
).all()
assert len(metadata) == 0
bindings = db_session_with_containers.scalars(
select(DatasetMetadataBinding).where(DatasetMetadataBinding.dataset_id == dataset.id)
).all()
assert len(bindings) == 0
# Check that process rules and queries were cleaned up
process_rules = db_session_with_containers.scalars(
select(DatasetProcessRule).where(DatasetProcessRule.dataset_id == dataset.id)
).all()
assert len(process_rules) == 0
queries = db_session_with_containers.scalars(
select(DatasetQuery).where(DatasetQuery.dataset_id == dataset.id)
).all()
assert len(queries) == 0
# Check that app dataset joins were cleaned up
app_joins = db_session_with_containers.scalars(
select(AppDatasetJoin).where(AppDatasetJoin.dataset_id == dataset.id)
).all()
assert len(app_joins) == 0
# Verify index processor was called
mock_index_processor = mock_external_service_dependencies["index_processor"]
mock_index_processor.clean.assert_called_once()
# Verify storage was not called (no files to delete)
mock_storage = mock_external_service_dependencies["storage"]
mock_storage.delete.assert_not_called()
def test_clean_dataset_task_success_with_documents_and_segments(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test successful dataset cleanup with documents and segments.
This test verifies that the task can successfully:
1. Clean up vector database indexes
2. Delete multiple documents and segments
3. Handle document segments with image references
4. Clean up storage files associated with documents
5. Remove all dataset-related data completely
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
# Create multiple documents
documents = []
for i in range(3):
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
documents.append(document)
# Create segments for each document
segments = []
for document in documents:
segment = self._create_test_segment(db_session_with_containers, account, tenant, dataset, document)
segments.append(segment)
# Create upload files for documents
upload_files = []
upload_file_ids = []
for document in documents:
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
upload_files.append(upload_file)
upload_file_ids.append(upload_file.id)
# Update document with file reference
import json
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
db_session_with_containers.commit()
# Create dataset metadata and bindings
metadata = DatasetMetadata(
dataset_id=dataset.id,
tenant_id=tenant.id,
name="test_metadata",
type=DatasetMetadataType.STRING,
created_by=account.id,
)
metadata.id = str(uuid.uuid4())
metadata.created_at = datetime.now()
binding = DatasetMetadataBinding(
tenant_id=tenant.id,
dataset_id=dataset.id,
metadata_id=metadata.id,
document_id=documents[0].id, # Use first document as example
created_by=account.id,
)
binding.id = str(uuid.uuid4())
binding.created_at = datetime.now()
db_session_with_containers.add(metadata)
db_session_with_containers.add(binding)
db_session_with_containers.commit()
# Execute the task
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
)
# Verify results
# Check that all documents were deleted
remaining_documents = db_session_with_containers.scalars(
select(Document).where(Document.dataset_id == dataset.id)
).all()
assert len(remaining_documents) == 0
# Check that all segments were deleted
remaining_segments = db_session_with_containers.scalars(
select(DocumentSegment).where(DocumentSegment.dataset_id == dataset.id)
).all()
assert len(remaining_segments) == 0
# Check that all upload files were deleted
remaining_files = db_session_with_containers.scalars(
select(UploadFile).where(UploadFile.id.in_(upload_file_ids))
).all()
assert len(remaining_files) == 0
# Check that metadata and bindings were cleaned up
remaining_metadata = db_session_with_containers.scalars(
select(DatasetMetadata).where(DatasetMetadata.dataset_id == dataset.id)
).all()
assert len(remaining_metadata) == 0
remaining_bindings = db_session_with_containers.scalars(
select(DatasetMetadataBinding).where(DatasetMetadataBinding.dataset_id == dataset.id)
).all()
assert len(remaining_bindings) == 0
# Verify index processor was called
mock_index_processor = mock_external_service_dependencies["index_processor"]
mock_index_processor.clean.assert_called_once()
# Verify storage delete was called for each file
mock_storage = mock_external_service_dependencies["storage"]
assert mock_storage.delete.call_count == 3
def test_clean_dataset_task_success_with_invalid_doc_form(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test successful dataset cleanup with invalid doc_form handling.
This test verifies that the task can successfully:
1. Handle None, empty, or whitespace-only doc_form values
2. Use default paragraph index type for cleanup
3. Continue with vector database cleanup using default type
4. Complete all cleanup operations successfully
5. Log appropriate warnings for invalid doc_form values
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
# Create a document and segment
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
segment = self._create_test_segment(db_session_with_containers, account, tenant, dataset, document)
# Execute the task with invalid doc_form values
test_cases = [None, "", " ", "\t\n"]
for invalid_doc_form in test_cases:
# Reset mock to clear previous calls
mock_index_processor = mock_external_service_dependencies["index_processor"]
mock_index_processor.clean.reset_mock()
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=invalid_doc_form,
)
# Verify that index processor was called with default type
mock_index_processor.clean.assert_called_once()
# Check that all data was cleaned up
remaining_documents = db_session_with_containers.scalars(
select(Document).where(Document.dataset_id == dataset.id)
).all()
assert len(remaining_documents) == 0
remaining_segments = db_session_with_containers.scalars(
select(DocumentSegment).where(DocumentSegment.dataset_id == dataset.id)
).all()
assert len(remaining_segments) == 0
# Recreate data for next test case
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
segment = self._create_test_segment(db_session_with_containers, account, tenant, dataset, document)
# Verify that IndexProcessorFactory was called with default type
mock_factory = mock_external_service_dependencies["index_processor_factory"]
# Should be called 4 times (once for each test case)
assert mock_factory.call_count == 4
def test_clean_dataset_task_error_handling_and_rollback(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test error handling and rollback mechanism when database operations fail.
This test verifies that the task can properly:
1. Handle database operation failures gracefully
2. Rollback database session to prevent dirty state
3. Continue cleanup operations even if some parts fail
4. Log appropriate error messages
5. Maintain database session integrity
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
segment = self._create_test_segment(db_session_with_containers, account, tenant, dataset, document)
# Mock IndexProcessorFactory to raise an exception
mock_index_processor = mock_external_service_dependencies["index_processor"]
mock_index_processor.clean.side_effect = Exception("Vector database cleanup failed")
# Execute the task - it should handle the exception gracefully
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
)
# Verify results - even with vector cleanup failure, documents and segments should be deleted
# Check that documents were still deleted despite vector cleanup failure
remaining_documents = db_session_with_containers.scalars(
select(Document).where(Document.dataset_id == dataset.id)
).all()
assert len(remaining_documents) == 0
# Check that segments were still deleted despite vector cleanup failure
remaining_segments = db_session_with_containers.scalars(
select(DocumentSegment).where(DocumentSegment.dataset_id == dataset.id)
).all()
assert len(remaining_segments) == 0
# Verify that index processor was called and failed
mock_index_processor.clean.assert_called_once()
# Verify that the task continued with cleanup despite the error
# This demonstrates the resilience of the cleanup process
def test_clean_dataset_task_with_image_file_references(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test dataset cleanup with image file references in document segments.
This test verifies that the task can properly:
1. Identify image upload file references in segment content
2. Clean up image files from storage
3. Remove image file database records
4. Handle multiple image references in segments
5. Clean up all image-related data completely
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
# Create image upload files
image_files = []
for i in range(3):
image_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
image_file.extension = ".jpg"
image_file.mime_type = "image/jpeg"
image_file.name = f"test_image_{i}.jpg"
image_files.append(image_file)
# Create segment with image references in content
segment_content = f"""
This is a test segment with image references.
"""
segment = DocumentSegment(
tenant_id=tenant.id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content=segment_content,
word_count=len(segment_content),
tokens=50,
created_by=account.id,
status=SegmentStatus.COMPLETED,
index_node_id=str(uuid.uuid4()),
index_node_hash="test_hash",
)
db_session_with_containers.add(segment)
db_session_with_containers.commit()
# Mock the get_image_upload_file_ids function to return our image file IDs
with patch("tasks.clean_dataset_task.get_image_upload_file_ids", autospec=True) as mock_get_image_ids:
mock_get_image_ids.return_value = [f.id for f in image_files]
# Execute the task
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
)
# Verify results
# Check that all documents were deleted
remaining_documents = db_session_with_containers.scalars(
select(Document).where(Document.dataset_id == dataset.id)
).all()
assert len(remaining_documents) == 0
# Check that all segments were deleted
remaining_segments = db_session_with_containers.scalars(
select(DocumentSegment).where(DocumentSegment.dataset_id == dataset.id)
).all()
assert len(remaining_segments) == 0
# Check that all image files were deleted from database
image_file_ids = [f.id for f in image_files]
remaining_image_files = db_session_with_containers.scalars(
select(UploadFile).where(UploadFile.id.in_(image_file_ids))
).all()
assert len(remaining_image_files) == 0
# Verify that storage.delete was called for each image file
mock_storage = mock_external_service_dependencies["storage"]
assert mock_storage.delete.call_count == 3
# Verify that get_image_upload_file_ids was called
mock_get_image_ids.assert_called_once()
def test_clean_dataset_task_performance_with_large_dataset(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test dataset cleanup performance with large amounts of data.
This test verifies that the task can efficiently:
1. Handle large numbers of documents and segments
2. Process multiple upload files efficiently
3. Maintain reasonable performance with complex data structures
4. Scale cleanup operations appropriately
5. Complete cleanup within acceptable time limits
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
# Create a large number of documents (simulating real-world scenario)
documents = []
segments = []
upload_files = []
upload_file_ids = []
# Create 50 documents with segments and upload files
for i in range(50):
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
documents.append(document)
# Create 3 segments per document
for j in range(3):
segment = self._create_test_segment(db_session_with_containers, account, tenant, dataset, document)
segments.append(segment)
# Create upload file for each document
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
upload_files.append(upload_file)
upload_file_ids.append(upload_file.id)
# Update document with file reference
import json
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
# Create dataset metadata and bindings
metadata_items = []
bindings = []
for i in range(10): # Create 10 metadata items
metadata = DatasetMetadata(
dataset_id=dataset.id,
tenant_id=tenant.id,
name=f"test_metadata_{i}",
type=DatasetMetadataType.STRING,
created_by=account.id,
)
metadata_items.append(metadata)
# Create binding for each metadata item
binding = DatasetMetadataBinding(
tenant_id=tenant.id,
dataset_id=dataset.id,
metadata_id=metadata.id,
document_id=documents[i % len(documents)].id,
created_by=account.id,
)
bindings.append(binding)
db_session_with_containers.add_all(metadata_items)
db_session_with_containers.add_all(bindings)
db_session_with_containers.commit()
# Measure cleanup performance
import time
start_time = time.time()
# Execute the task
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
)
end_time = time.time()
cleanup_duration = end_time - start_time
# Verify results
# Check that all documents were deleted
remaining_documents = db_session_with_containers.scalars(
select(Document).where(Document.dataset_id == dataset.id)
).all()
assert len(remaining_documents) == 0
# Check that all segments were deleted
remaining_segments = db_session_with_containers.scalars(
select(DocumentSegment).where(DocumentSegment.dataset_id == dataset.id)
).all()
assert len(remaining_segments) == 0
# Check that all upload files were deleted
remaining_files = db_session_with_containers.scalars(
select(UploadFile).where(UploadFile.id.in_(upload_file_ids))
).all()
assert len(remaining_files) == 0
# Check that all metadata and bindings were deleted
remaining_metadata = db_session_with_containers.scalars(
select(DatasetMetadata).where(DatasetMetadata.dataset_id == dataset.id)
).all()
assert len(remaining_metadata) == 0
remaining_bindings = db_session_with_containers.scalars(
select(DatasetMetadataBinding).where(DatasetMetadataBinding.dataset_id == dataset.id)
).all()
assert len(remaining_bindings) == 0
# Verify performance expectations
# Cleanup should complete within reasonable time (adjust threshold as needed)
assert cleanup_duration < 10.0, f"Cleanup took too long: {cleanup_duration:.2f} seconds"
# Verify that storage.delete was called for each file
mock_storage = mock_external_service_dependencies["storage"]
assert mock_storage.delete.call_count == 50
# Verify that index processor was called
mock_index_processor = mock_external_service_dependencies["index_processor"]
mock_index_processor.clean.assert_called_once()
# Log performance metrics
print("\nPerformance Test Results:")
print(f"Documents processed: {len(documents)}")
print(f"Segments processed: {len(segments)}")
print(f"Upload files processed: {len(upload_files)}")
print(f"Metadata items processed: {len(metadata_items)}")
print(f"Total cleanup time: {cleanup_duration:.3f} seconds")
print(f"Average time per document: {cleanup_duration / len(documents):.3f} seconds")
def test_clean_dataset_task_storage_exception_handling(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test dataset cleanup when storage operations fail.
This test verifies that the task can properly:
1. Handle storage deletion failures gracefully
2. Continue cleanup process despite storage errors
3. Log appropriate error messages for storage failures
4. Maintain database consistency even with storage issues
5. Provide meaningful error reporting
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
segment = self._create_test_segment(db_session_with_containers, account, tenant, dataset, document)
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
# Update document with file reference
import json
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
db_session_with_containers.commit()
# Mock storage to raise exceptions
mock_storage = mock_external_service_dependencies["storage"]
mock_storage.delete.side_effect = Exception("Storage service unavailable")
# Execute the task - it should handle storage failures gracefully
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
)
# Verify results
# Note: When storage operations fail, database deletions may be rolled back by implementation.
# This test focuses on ensuring the task handles the exception and continues execution/logging.
# Check that upload file was still deleted from database despite storage failure
# Note: When storage operations fail, the upload file may not be deleted
# This demonstrates that the cleanup process continues even with storage errors
remaining_files = db_session_with_containers.scalars(
select(UploadFile).where(UploadFile.id == upload_file.id)
).all()
# The upload file should still be deleted from the database even if storage cleanup fails
# However, this depends on the specific implementation of clean_dataset_task
if len(remaining_files) > 0:
print(f"Warning: Upload file {upload_file.id} was not deleted despite storage failure")
print("This demonstrates that the cleanup process continues even with storage errors")
# We don't assert here as the behavior depends on the specific implementation
# Verify that storage.delete was called
mock_storage.delete.assert_called_once()
# Verify that index processor was called successfully
mock_index_processor = mock_external_service_dependencies["index_processor"]
mock_index_processor.clean.assert_called_once()
# This test demonstrates that the cleanup process continues
# even when external storage operations fail, ensuring data
# consistency in the database
def test_clean_dataset_task_edge_cases_and_boundary_conditions(
self, db_session_with_containers: Session, mock_external_service_dependencies
):
"""
Test dataset cleanup with edge cases and boundary conditions.
This test verifies that the task can properly:
1. Handle datasets with no documents or segments
2. Process datasets with minimal metadata
3. Handle extremely long dataset names and descriptions
4. Process datasets with special characters in content
5. Handle datasets with maximum allowed field values
"""
# Create test data with edge cases
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
# Create dataset with long name and description (within database limits)
long_name = "a" * 250 # Long name within varchar(255) limit
long_description = "b" * 500 # Long description within database limits
dataset = Dataset(
id=str(uuid.uuid4()),
tenant_id=tenant.id,
name=long_name,
description=long_description,
indexing_technique=IndexTechniqueType.HIGH_QUALITY,
index_struct='{"type": "paragraph", "max_length": 10000}',
collection_binding_id=str(uuid.uuid4()),
created_by=account.id,
created_at=datetime.now(),
updated_at=datetime.now(),
)
db_session_with_containers.add(dataset)
db_session_with_containers.commit()
# Create document with special characters in name
special_content = "Special chars: !@#$%^&*()_+-=[]{}|;':\",./<>?`~"
document = Document(
id=str(uuid.uuid4()),
tenant_id=tenant.id,
dataset_id=dataset.id,
position=1,
data_source_type=DataSourceType.UPLOAD_FILE,
data_source_info="{}",
batch="test_batch",
name=f"test_doc_{special_content}",
created_from=DocumentCreatedFrom.WEB,
created_by=account.id,
created_at=datetime.now(),
updated_at=datetime.now(),
)
db_session_with_containers.add(document)
db_session_with_containers.commit()
# Create segment with special characters and very long content
long_content = "Very long content " * 100 # Long content within reasonable limits
segment_content = f"Segment with special chars: {special_content}\n{long_content}"
segment = DocumentSegment(
tenant_id=tenant.id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content=segment_content,
word_count=len(segment_content.split()),
tokens=len(segment_content) // 4, # Rough token estimation
created_by=account.id,
status=SegmentStatus.COMPLETED,
index_node_id=str(uuid.uuid4()),
index_node_hash="test_hash_" + "x" * 50, # Long hash within limits
)
db_session_with_containers.add(segment)
db_session_with_containers.commit()
# Create upload file with special characters in name
special_filename = f"test_file_{special_content}.txt"
upload_file = UploadFile(
tenant_id=tenant.id,
storage_type=StorageType.LOCAL,
key=f"test_files/{special_filename}",
name=special_filename,
size=1024,
extension=".txt",
mime_type="text/plain",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(),
used=False,
)
db_session_with_containers.add(upload_file)
db_session_with_containers.commit()
# Update document with file reference
import json
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
db_session_with_containers.commit()
# Save upload file ID for verification
upload_file_id = upload_file.id
# Create metadata with special characters
special_metadata = DatasetMetadata(
dataset_id=dataset.id,
tenant_id=tenant.id,
name=f"metadata_{special_content}",
type=DatasetMetadataType.STRING,
created_by=account.id,
)
special_metadata.id = str(uuid.uuid4())
special_metadata.created_at = datetime.now()
db_session_with_containers.add(special_metadata)
db_session_with_containers.commit()
# Execute the task
clean_dataset_task(
dataset_id=dataset.id,
tenant_id=tenant.id,
indexing_technique=dataset.indexing_technique,
index_struct=dataset.index_struct,
collection_binding_id=dataset.collection_binding_id,
doc_form=dataset.doc_form,
)
# Verify results
# Check that all documents were deleted
remaining_documents = db_session_with_containers.scalars(
select(Document).where(Document.dataset_id == dataset.id)
).all()
assert len(remaining_documents) == 0
# Check that all segments were deleted
remaining_segments = db_session_with_containers.scalars(
select(DocumentSegment).where(DocumentSegment.dataset_id == dataset.id)
).all()
assert len(remaining_segments) == 0
# Check that all upload files were deleted
remaining_files = db_session_with_containers.scalars(
select(UploadFile).where(UploadFile.id == upload_file_id)
).all()
assert len(remaining_files) == 0
# Check that all metadata was deleted
remaining_metadata = db_session_with_containers.scalars(
select(DatasetMetadata).where(DatasetMetadata.dataset_id == dataset.id)
).all()
assert len(remaining_metadata) == 0
# Verify that storage.delete was called
mock_storage = mock_external_service_dependencies["storage"]
mock_storage.delete.assert_called_once()
# Verify that index processor was called
mock_index_processor = mock_external_service_dependencies["index_processor"]
mock_index_processor.clean.assert_called_once()
# This test demonstrates that the cleanup process can handle
# extreme edge cases including very long content, special characters,
# and boundary conditions without failing