mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 02:28:30 +08:00
test: use sqlite3 session in test_dataset_service_lock_not_owned (#38755)
This commit is contained in:
Asuka Minato
GitHub
parent
c12d4662db
commit
713c577b51
@ -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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Loading…
Reference in New Issue
Block a user