""" 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 src.web.services.autolabel import AutoLabelService from src.data.admin_db import AdminDB 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 MockAdminDB: """Mock AdminDB for testing.""" def __init__(self): self.documents = {} self.annotations = [] self.status_updates = [] def get_document(self, document_id): """Get document by ID.""" return self.documents.get(str(document_id)) def update_document_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 def delete_annotations_for_document(self, document_id, source=None): """Mock delete annotations.""" return 0 def create_annotations_batch(self, annotations): """Mock create annotations.""" self.annotations.extend(annotations) @pytest.fixture def mock_db(): """Create mock admin DB.""" return MockAdminDB() @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, db, page_number=1): return 0 # No annotations created (mocked) monkeypatch.setattr(AutoLabelService, "_process_image", mock_process_image) return service class TestAutoLabelWithLocks: """Tests for auto-label service with lock integration.""" def test_auto_label_unlocked_document_succeeds(self, auto_label_service, mock_db, tmp_path): """Test auto-labeling succeeds on unlocked document.""" # Create test document (unlocked) document_id = str(uuid4()) mock_db.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"}, db=mock_db, ) # Should succeed assert result["status"] == "completed" # Verify status was updated to running and then completed assert len(mock_db.status_updates) >= 2 assert mock_db.status_updates[0]["auto_label_status"] == "running" def test_auto_label_locked_document_fails(self, auto_label_service, mock_db, 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_db.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"}, db=mock_db, ) # 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_db.status_updates ) def test_auto_label_expired_lock_succeeds(self, auto_label_service, mock_db, 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_db.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"}, db=mock_db, ) # Should succeed (lock expired) assert result["status"] == "completed" def test_auto_label_skip_lock_check(self, auto_label_service, mock_db, 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_db.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"}, db=mock_db, 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_db, 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"}, db=mock_db, ) # 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_db, 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_db.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"}, db=mock_db, # skip_lock_check not specified, should default to False ) # Should fail due to lock assert result["status"] == "failed" assert "locked" in result["error"].lower()