272 lines
9.2 KiB
Python
272 lines
9.2 KiB
Python
"""
|
|
Tests for Auto-Label Service with Annotation Lock Integration (Phase 3.5).
|
|
"""
|
|
|
|
import pytest
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, MagicMock
|
|
from uuid import uuid4
|
|
|
|
from inference.web.services.autolabel import AutoLabelService
|
|
|
|
|
|
class MockDocument:
|
|
"""Mock document for testing."""
|
|
|
|
def __init__(self, document_id, annotation_lock_until=None):
|
|
self.document_id = document_id
|
|
self.annotation_lock_until = annotation_lock_until
|
|
self.status = "pending"
|
|
self.auto_label_status = None
|
|
self.auto_label_error = None
|
|
|
|
|
|
class MockDocumentRepository:
|
|
"""Mock DocumentRepository for testing."""
|
|
|
|
def __init__(self):
|
|
self.documents = {}
|
|
self.status_updates = []
|
|
|
|
def get(self, document_id):
|
|
"""Get document by ID."""
|
|
return self.documents.get(str(document_id))
|
|
|
|
def update_status(
|
|
self,
|
|
document_id,
|
|
status=None,
|
|
auto_label_status=None,
|
|
auto_label_error=None,
|
|
):
|
|
"""Mock status update."""
|
|
self.status_updates.append({
|
|
"document_id": document_id,
|
|
"status": status,
|
|
"auto_label_status": auto_label_status,
|
|
"auto_label_error": auto_label_error,
|
|
})
|
|
doc = self.documents.get(str(document_id))
|
|
if doc:
|
|
if status:
|
|
doc.status = status
|
|
if auto_label_status:
|
|
doc.auto_label_status = auto_label_status
|
|
if auto_label_error:
|
|
doc.auto_label_error = auto_label_error
|
|
|
|
|
|
class MockAnnotationRepository:
|
|
"""Mock AnnotationRepository for testing."""
|
|
|
|
def __init__(self):
|
|
self.annotations = []
|
|
|
|
def delete_for_document(self, document_id, source=None):
|
|
"""Mock delete annotations."""
|
|
return 0
|
|
|
|
def create_batch(self, annotations):
|
|
"""Mock create annotations."""
|
|
self.annotations.extend(annotations)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_doc_repo():
|
|
"""Create mock document repository."""
|
|
return MockDocumentRepository()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_ann_repo():
|
|
"""Create mock annotation repository."""
|
|
return MockAnnotationRepository()
|
|
|
|
|
|
@pytest.fixture
|
|
def auto_label_service(monkeypatch):
|
|
"""Create auto-label service with mocked image processing."""
|
|
service = AutoLabelService()
|
|
# Mock the OCR engine to avoid dependencies
|
|
service._ocr_engine = Mock()
|
|
service._ocr_engine.extract_from_image = Mock(return_value=[])
|
|
|
|
# Mock the image processing methods to avoid file I/O errors
|
|
def mock_process_image(self, document_id, image_path, field_values, ann_repo, page_number=1):
|
|
return 0 # No annotations created (mocked)
|
|
|
|
def mock_process_pdf(self, document_id, pdf_path, field_values, ann_repo):
|
|
return 0 # No annotations created (mocked)
|
|
|
|
monkeypatch.setattr(AutoLabelService, "_process_image", mock_process_image)
|
|
monkeypatch.setattr(AutoLabelService, "_process_pdf", mock_process_pdf)
|
|
|
|
return service
|
|
|
|
|
|
class TestAutoLabelWithLocks:
|
|
"""Tests for auto-label service with lock integration."""
|
|
|
|
def test_auto_label_unlocked_document_succeeds(self, auto_label_service, mock_doc_repo, mock_ann_repo, tmp_path):
|
|
"""Test auto-labeling succeeds on unlocked document."""
|
|
# Create test document (unlocked)
|
|
document_id = str(uuid4())
|
|
mock_doc_repo.documents[document_id] = MockDocument(
|
|
document_id=document_id,
|
|
annotation_lock_until=None,
|
|
)
|
|
|
|
# Create dummy file
|
|
test_file = tmp_path / "test.png"
|
|
test_file.write_text("dummy")
|
|
|
|
# Attempt auto-label
|
|
result = auto_label_service.auto_label_document(
|
|
document_id=document_id,
|
|
file_path=str(test_file),
|
|
field_values={"invoice_number": "INV-001"},
|
|
doc_repo=mock_doc_repo,
|
|
ann_repo=mock_ann_repo,
|
|
)
|
|
|
|
# Should succeed
|
|
assert result["status"] == "completed"
|
|
# Verify status was updated to running and then completed
|
|
assert len(mock_doc_repo.status_updates) >= 2
|
|
assert mock_doc_repo.status_updates[0]["auto_label_status"] == "running"
|
|
|
|
def test_auto_label_locked_document_fails(self, auto_label_service, mock_doc_repo, mock_ann_repo, tmp_path):
|
|
"""Test auto-labeling fails on locked document."""
|
|
# Create test document (locked for 1 hour)
|
|
document_id = str(uuid4())
|
|
lock_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
|
mock_doc_repo.documents[document_id] = MockDocument(
|
|
document_id=document_id,
|
|
annotation_lock_until=lock_until,
|
|
)
|
|
|
|
# Create dummy file
|
|
test_file = tmp_path / "test.png"
|
|
test_file.write_text("dummy")
|
|
|
|
# Attempt auto-label (should fail)
|
|
result = auto_label_service.auto_label_document(
|
|
document_id=document_id,
|
|
file_path=str(test_file),
|
|
field_values={"invoice_number": "INV-001"},
|
|
doc_repo=mock_doc_repo,
|
|
ann_repo=mock_ann_repo,
|
|
)
|
|
|
|
# Should fail
|
|
assert result["status"] == "failed"
|
|
assert "locked for annotation" in result["error"]
|
|
assert result["annotations_created"] == 0
|
|
|
|
# Verify status was updated to failed
|
|
assert any(
|
|
update["auto_label_status"] == "failed"
|
|
for update in mock_doc_repo.status_updates
|
|
)
|
|
|
|
def test_auto_label_expired_lock_succeeds(self, auto_label_service, mock_doc_repo, mock_ann_repo, tmp_path):
|
|
"""Test auto-labeling succeeds when lock has expired."""
|
|
# Create test document (lock expired 1 hour ago)
|
|
document_id = str(uuid4())
|
|
lock_until = datetime.now(timezone.utc) - timedelta(hours=1)
|
|
mock_doc_repo.documents[document_id] = MockDocument(
|
|
document_id=document_id,
|
|
annotation_lock_until=lock_until,
|
|
)
|
|
|
|
# Create dummy file
|
|
test_file = tmp_path / "test.png"
|
|
test_file.write_text("dummy")
|
|
|
|
# Attempt auto-label
|
|
result = auto_label_service.auto_label_document(
|
|
document_id=document_id,
|
|
file_path=str(test_file),
|
|
field_values={"invoice_number": "INV-001"},
|
|
doc_repo=mock_doc_repo,
|
|
ann_repo=mock_ann_repo,
|
|
)
|
|
|
|
# Should succeed (lock expired)
|
|
assert result["status"] == "completed"
|
|
|
|
def test_auto_label_skip_lock_check(self, auto_label_service, mock_doc_repo, mock_ann_repo, tmp_path):
|
|
"""Test auto-labeling with skip_lock_check=True bypasses lock."""
|
|
# Create test document (locked)
|
|
document_id = str(uuid4())
|
|
lock_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
|
mock_doc_repo.documents[document_id] = MockDocument(
|
|
document_id=document_id,
|
|
annotation_lock_until=lock_until,
|
|
)
|
|
|
|
# Create dummy file
|
|
test_file = tmp_path / "test.png"
|
|
test_file.write_text("dummy")
|
|
|
|
# Attempt auto-label with skip_lock_check=True
|
|
result = auto_label_service.auto_label_document(
|
|
document_id=document_id,
|
|
file_path=str(test_file),
|
|
field_values={"invoice_number": "INV-001"},
|
|
doc_repo=mock_doc_repo,
|
|
ann_repo=mock_ann_repo,
|
|
skip_lock_check=True, # Bypass lock check
|
|
)
|
|
|
|
# Should succeed even though document is locked
|
|
assert result["status"] == "completed"
|
|
|
|
def test_auto_label_document_not_found(self, auto_label_service, mock_doc_repo, mock_ann_repo, tmp_path):
|
|
"""Test auto-labeling fails when document doesn't exist."""
|
|
# Create dummy file
|
|
test_file = tmp_path / "test.png"
|
|
test_file.write_text("dummy")
|
|
|
|
# Attempt auto-label on non-existent document
|
|
result = auto_label_service.auto_label_document(
|
|
document_id=str(uuid4()),
|
|
file_path=str(test_file),
|
|
field_values={"invoice_number": "INV-001"},
|
|
doc_repo=mock_doc_repo,
|
|
ann_repo=mock_ann_repo,
|
|
)
|
|
|
|
# Should fail
|
|
assert result["status"] == "failed"
|
|
assert "not found" in result["error"]
|
|
|
|
def test_auto_label_respects_lock_by_default(self, auto_label_service, mock_doc_repo, mock_ann_repo, tmp_path):
|
|
"""Test that lock check is enabled by default."""
|
|
# Create test document (locked)
|
|
document_id = str(uuid4())
|
|
lock_until = datetime.now(timezone.utc) + timedelta(minutes=30)
|
|
mock_doc_repo.documents[document_id] = MockDocument(
|
|
document_id=document_id,
|
|
annotation_lock_until=lock_until,
|
|
)
|
|
|
|
# Create dummy file
|
|
test_file = tmp_path / "test.png"
|
|
test_file.write_text("dummy")
|
|
|
|
# Call without explicit skip_lock_check (defaults to False)
|
|
result = auto_label_service.auto_label_document(
|
|
document_id=document_id,
|
|
file_path=str(test_file),
|
|
field_values={"invoice_number": "INV-001"},
|
|
doc_repo=mock_doc_repo,
|
|
ann_repo=mock_ann_repo,
|
|
# skip_lock_check not specified, should default to False
|
|
)
|
|
|
|
# Should fail due to lock
|
|
assert result["status"] == "failed"
|
|
assert "locked" in result["error"].lower()
|