diff --git a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py index 5e5d406edb8..50ca483b976 100644 --- a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py +++ b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py @@ -3,12 +3,20 @@ from unittest.mock import Mock, create_autospec import pytest from redis.exceptions import LockNotOwnedError +from sqlalchemy import func, select +from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from models.account import Account -from models.dataset import Dataset, Document +from models.account import Account, Tenant +from models.dataset import Dataset, DatasetProcessRule, Document, DocumentSegment +from models.enums import ProcessRuleMode from services.dataset_service import DocumentService, SegmentService +TENANT_ID = "11111111-1111-1111-1111-111111111111" +USER_ID = "22222222-2222-2222-2222-222222222222" +DATASET_ID = "33333333-3333-3333-3333-333333333333" +DOCUMENT_ID = "44444444-4444-4444-4444-444444444444" + class FakeLock: """Lock that always fails on enter with LockNotOwnedError.""" @@ -23,9 +31,11 @@ class FakeLock: @pytest.fixture def fake_current_user(monkeypatch: pytest.MonkeyPatch): - user = create_autospec(Account, instance=True) - user.id = "user-1" - user.current_tenant_id = "tenant-1" + tenant = Tenant(name="Test Tenant") + tenant.id = TENANT_ID + user = Account(name="Test User", email="test@example.com") + user.id = USER_ID + user._current_tenant = tenant monkeypatch.setattr("services.dataset_service.current_user", user) return user @@ -60,18 +70,32 @@ def fake_lock(monkeypatch: pytest.MonkeyPatch): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, Dataset, DatasetProcessRule, Document)], indirect=True) def test_save_document_with_dataset_id_ignores_lock_not_owned( monkeypatch: pytest.MonkeyPatch, fake_current_user, fake_features, fake_lock, + sqlite_session: Session, ): # Arrange - dataset = create_autospec(Dataset, instance=True) - dataset.id = "ds-1" - dataset.tenant_id = fake_current_user.current_tenant_id - dataset.data_source_type = "upload_file" - dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY # so we skip re-initialization branch + dataset = Dataset( + id=DATASET_ID, + tenant_id=TENANT_ID, + name="Test Dataset", + description="", + created_by=USER_ID, + data_source_type="upload_file", + indexing_technique=IndexTechniqueType.HIGH_QUALITY, + ) + process_rule = DatasetProcessRule( + dataset_id=DATASET_ID, + mode=ProcessRuleMode.AUTOMATIC, + rules=None, + created_by=USER_ID, + ) + sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, process_rule]) + sqlite_session.commit() # Minimal knowledge_config stub that satisfies pre-lock code info_list = types.SimpleNamespace(data_source_type="upload_file") @@ -93,10 +117,6 @@ def test_save_document_with_dataset_id_ignores_lock_not_owned( # Avoid touching real doc_form logic monkeypatch.setattr("services.dataset_service.DatasetService.check_doc_form", lambda *a, **k: None) - # Avoid real DB interactions - db_mock = Mock() - db_mock.session = Mock() - monkeypatch.setattr("services.dataset_service.db", db_mock) # Act: this would hit the redis lock, whose __enter__ raises LockNotOwnedError. # Our implementation should catch it and still return (documents, batch). @@ -104,15 +124,19 @@ def test_save_document_with_dataset_id_ignores_lock_not_owned( dataset=dataset, knowledge_config=knowledge_config, account=account, - session=db_mock.session, + dataset_process_rule=process_rule, + session=sqlite_session, ) # Assert # We mainly care that: # - No exception is raised # - The function returns a sensible tuple - assert isinstance(documents, list) - assert isinstance(batch, str) + assert documents == [] + assert batch + assert not sqlite_session.in_transaction() + assert sqlite_session.scalar(select(func.count(Document.id))) == 0 + assert sqlite_session.get(DatasetProcessRule, process_rule.id) is process_rule # --------------------------------------------------------------------------- @@ -120,22 +144,38 @@ def test_save_document_with_dataset_id_ignores_lock_not_owned( # --------------------------------------------------------------------------- +@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, Dataset, Document, DocumentSegment)], indirect=True) def test_add_segment_ignores_lock_not_owned( monkeypatch: pytest.MonkeyPatch, fake_current_user, fake_lock, + sqlite_session: Session, ): # Arrange - dataset = create_autospec(Dataset, instance=True) - dataset.id = "ds-1" - dataset.tenant_id = fake_current_user.current_tenant_id - dataset.indexing_technique = IndexTechniqueType.ECONOMY # skip embedding/token calculation branch - - document = create_autospec(Document, instance=True) - document.id = "doc-1" - document.dataset_id = dataset.id - document.word_count = 0 - document.doc_form = IndexStructureType.QA_INDEX + dataset = Dataset( + id=DATASET_ID, + tenant_id=TENANT_ID, + name="Test Dataset", + description="", + created_by=USER_ID, + indexing_technique=IndexTechniqueType.ECONOMY, + ) + document = Document( + id=DOCUMENT_ID, + tenant_id=TENANT_ID, + dataset_id=DATASET_ID, + position=1, + data_source_type="upload_file", + data_source_info="{}", + batch="batch-1", + name="Test Document", + created_from="web", + created_by=USER_ID, + word_count=0, + doc_form=IndexStructureType.QA_INDEX, + ) + sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, document]) + sqlite_session.commit() # Minimal args required by add_segment args = { @@ -144,18 +184,18 @@ def test_add_segment_ignores_lock_not_owned( "keywords": ["k1", "k2"], } - # Avoid real DB operations - db_mock = Mock() - db_mock.session = Mock() - monkeypatch.setattr("services.dataset_service.db", db_mock) monkeypatch.setattr("services.dataset_service.VectorService", Mock()) # Act - result = SegmentService.create_segment(args=args, document=document, dataset=dataset, session=db_mock.session) + result = SegmentService.create_segment(args=args, document=document, dataset=dataset, session=sqlite_session) # Assert # Under LockNotOwnedError except, add_segment should swallow the error and return None. assert result is None + assert not sqlite_session.in_transaction() + assert sqlite_session.scalar(select(func.count(DocumentSegment.id))) == 0 + sqlite_session.refresh(document) + assert document.word_count == 0 # ---------------------------------------------------------------------------