WIP

tests/web/test_annotation_phase5.py

@@ -9,8 +9,12 @@ from uuid import uuid4
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
-from inference.web.api.v1.admin.annotations import create_annotation_router
-from inference.web.core.auth import validate_admin_token, get_admin_db
+from inference.web.api.v1.admin.annotations import (
+    create_annotation_router,
+    get_doc_repository,
+    get_ann_repository,
+)
+from inference.web.core.auth import validate_admin_token
 
 
 class MockAdminDocument:
@@ -73,22 +77,40 @@ class MockAnnotationHistory:
         self.created_at = kwargs.get('created_at', datetime.utcnow())
 
 
-class MockAdminDB:
-    """Mock AdminDB for testing Phase 5."""
+class MockDocumentRepository:
+    """Mock DocumentRepository for testing Phase 5."""
 
     def __init__(self):
         self.documents = {}
-        self.annotations = {}
-        self.annotation_history = {}
 
-    def get_document_by_token(self, document_id, admin_token):
+    def get(self, document_id):
+        """Get document by ID."""
+        return self.documents.get(str(document_id))
+
+    def get_by_token(self, document_id, admin_token=None):
         """Get document by ID and token."""
         doc = self.documents.get(str(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 verify_annotation(self, annotation_id, admin_token):
+
+class MockAnnotationRepository:
+    """Mock AnnotationRepository for testing Phase 5."""
+
+    def __init__(self):
+        self.annotations = {}
+        self.annotation_history = {}
+
+    def get(self, annotation_id):
+        """Get annotation by ID."""
+        return self.annotations.get(str(annotation_id))
+
+    def get_for_document(self, document_id, page_number=None):
+        """Get annotations for a document."""
+        return [a for a in self.annotations.values() if str(a.document_id) == str(document_id)]
+
+    def verify(self, annotation_id, admin_token):
         """Mark annotation as verified."""
         annotation = self.annotations.get(str(annotation_id))
         if annotation:
@@ -98,7 +120,7 @@ class MockAdminDB:
             return annotation
         return None
 
-    def override_annotation(
+    def override(
         self,
         annotation_id,
         admin_token,
@@ -131,7 +153,7 @@ class MockAdminDB:
             return annotation
         return None
 
-    def get_annotation_history(self, annotation_id):
+    def get_history(self, annotation_id):
         """Get annotation history."""
         return self.annotation_history.get(str(annotation_id), [])
 
@@ -141,15 +163,16 @@ def app():
     """Create test FastAPI app."""
     app = FastAPI()
 
-    # Create mock DB
-    mock_db = MockAdminDB()
+    # Create mock repositories
+    mock_document_repo = MockDocumentRepository()
+    mock_annotation_repo = MockAnnotationRepository()
 
     # Add test document
     doc1 = MockAdminDocument(
         filename="TEST001.pdf",
         status="labeled",
     )
-    mock_db.documents[str(doc1.document_id)] = doc1
+    mock_document_repo.documents[str(doc1.document_id)] = doc1
 
     # Add test annotations
     ann1 = MockAnnotation(
@@ -169,8 +192,8 @@ def app():
         confidence=0.98,
     )
 
-    mock_db.annotations[str(ann1.annotation_id)] = ann1
-    mock_db.annotations[str(ann2.annotation_id)] = ann2
+    mock_annotation_repo.annotations[str(ann1.annotation_id)] = ann1
+    mock_annotation_repo.annotations[str(ann2.annotation_id)] = ann2
 
     # Store document ID and annotation IDs for tests
     app.state.document_id = str(doc1.document_id)
@@ -179,7 +202,8 @@ def app():
 
     # Override dependencies
     app.dependency_overrides[validate_admin_token] = lambda: "test-token"
-    app.dependency_overrides[get_admin_db] = lambda: mock_db
+    app.dependency_overrides[get_doc_repository] = lambda: mock_document_repo
+    app.dependency_overrides[get_ann_repository] = lambda: mock_annotation_repo
 
     # Include router
     router = create_annotation_router()