WIP

tests/web/test_annotation_locks.py

@@ -10,7 +10,10 @@ from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
 from inference.web.api.v1.admin.locks import create_locks_router
-from inference.web.core.auth import validate_admin_token, get_admin_db
+from inference.web.core.auth import (
+    validate_admin_token,
+    get_document_repository,
+)
 
 
 class MockAdminDocument:
@@ -34,23 +37,27 @@ class MockAdminDocument:
         self.updated_at = kwargs.get('updated_at', datetime.utcnow())
 
 
-class MockAdminDB:
-    """Mock AdminDB for testing annotation locks."""
+class MockDocumentRepository:
+    """Mock DocumentRepository for testing annotation locks."""
 
     def __init__(self):
         self.documents = {}
 
-    def get_document_by_token(self, document_id, admin_token):
+    def get(self, document_id):
+        """Get single document by ID."""
+        return self.documents.get(document_id)
+
+    def get_by_token(self, document_id, admin_token=None):
         """Get single document by ID and token."""
         doc = self.documents.get(document_id)
-        if doc and doc.admin_token == admin_token:
+        if doc and (admin_token is None or doc.admin_token == admin_token):
             return doc
         return None
 
-    def acquire_annotation_lock(self, document_id, admin_token, duration_seconds=300):
+    def acquire_annotation_lock(self, document_id, admin_token=None, duration_seconds=300):
         """Acquire annotation lock for a document."""
         doc = self.documents.get(document_id)
-        if not doc or doc.admin_token != admin_token:
+        if not doc:
             return None
 
         # Check if already locked
@@ -62,20 +69,20 @@ class MockAdminDB:
         doc.annotation_lock_until = now + timedelta(seconds=duration_seconds)
         return doc
 
-    def release_annotation_lock(self, document_id, admin_token, force=False):
+    def release_annotation_lock(self, document_id, admin_token=None, force=False):
         """Release annotation lock for a document."""
         doc = self.documents.get(document_id)
-        if not doc or doc.admin_token != admin_token:
+        if not doc:
             return None
 
         # Release lock
         doc.annotation_lock_until = None
         return doc
 
-    def extend_annotation_lock(self, document_id, admin_token, additional_seconds=300):
+    def extend_annotation_lock(self, document_id, admin_token=None, additional_seconds=300):
         """Extend an existing annotation lock."""
         doc = self.documents.get(document_id)
-        if not doc or doc.admin_token != admin_token:
+        if not doc:
             return None
 
         # Check if lock exists and is still valid
@@ -93,8 +100,8 @@ def app():
     """Create test FastAPI app."""
     app = FastAPI()
 
-    # Create mock DB
-    mock_db = MockAdminDB()
+    # Create mock repository
+    mock_document_repo = MockDocumentRepository()
 
     # Add test document
     doc1 = MockAdminDocument(
@@ -103,11 +110,11 @@ def app():
         upload_source="ui",
     )
 
-    mock_db.documents[str(doc1.document_id)] = doc1
+    mock_document_repo.documents[str(doc1.document_id)] = doc1
 
     # Override dependencies
     app.dependency_overrides[validate_admin_token] = lambda: "test-token"
-    app.dependency_overrides[get_admin_db] = lambda: mock_db
+    app.dependency_overrides[get_document_repository] = lambda: mock_document_repo
 
     # Include router
     router = create_locks_router()
@@ -124,9 +131,9 @@ def client(app):
 
 @pytest.fixture
 def document_id(app):
-    """Get document ID from the mock DB."""
-    mock_db = app.dependency_overrides[get_admin_db]()
-    return str(list(mock_db.documents.keys())[0])
+    """Get document ID from the mock repository."""
+    mock_document_repo = app.dependency_overrides[get_document_repository]()
+    return str(list(mock_document_repo.documents.keys())[0])
 
 
 class TestAnnotationLocks: