WIP

tests/web/test_autolabel_with_locks.py

@@ -0,0 +1,250 @@
+"""
+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()