WIP

tests/data/repositories/__init__.py

@@ -0,0 +1 @@
+"""Tests for repository pattern implementation."""

tests/data/repositories/test_annotation_repository.py

@@ -0,0 +1,711 @@
+"""
+Tests for AnnotationRepository
+
+100% coverage tests for annotation management.
+"""
+
+import pytest
+from datetime import datetime, timezone
+from unittest.mock import MagicMock, patch
+from uuid import uuid4, UUID
+
+from inference.data.admin_models import AdminAnnotation, AnnotationHistory
+from inference.data.repositories.annotation_repository import AnnotationRepository
+
+
+class TestAnnotationRepository:
+    """Tests for AnnotationRepository."""
+
+    @pytest.fixture
+    def sample_annotation(self) -> AdminAnnotation:
+        """Create a sample annotation for testing."""
+        return AdminAnnotation(
+            annotation_id=uuid4(),
+            document_id=uuid4(),
+            page_number=1,
+            class_id=0,
+            class_name="invoice_number",
+            x_center=0.5,
+            y_center=0.3,
+            width=0.2,
+            height=0.05,
+            bbox_x=100,
+            bbox_y=200,
+            bbox_width=150,
+            bbox_height=30,
+            text_value="INV-001",
+            confidence=0.95,
+            source="auto",
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def sample_history(self) -> AnnotationHistory:
+        """Create a sample annotation history for testing."""
+        return AnnotationHistory(
+            history_id=uuid4(),
+            annotation_id=uuid4(),
+            document_id=uuid4(),
+            action="override",
+            previous_value={"class_name": "old_class"},
+            new_value={"class_name": "new_class"},
+            changed_by="admin-token",
+            change_reason="Correction",
+            created_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def repo(self) -> AnnotationRepository:
+        """Create an AnnotationRepository instance."""
+        return AnnotationRepository()
+
+    # =========================================================================
+    # create() tests
+    # =========================================================================
+
+    def test_create_returns_annotation_id(self, repo):
+        """Test create returns annotation ID."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                document_id=str(uuid4()),
+                page_number=1,
+                class_id=0,
+                class_name="invoice_number",
+                x_center=0.5,
+                y_center=0.3,
+                width=0.2,
+                height=0.05,
+                bbox_x=100,
+                bbox_y=200,
+                bbox_width=150,
+                bbox_height=30,
+            )
+
+            assert result is not None
+            mock_session.add.assert_called_once()
+            mock_session.flush.assert_called_once()
+
+    def test_create_with_optional_params(self, repo):
+        """Test create with optional text_value and confidence."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                document_id=str(uuid4()),
+                page_number=2,
+                class_id=1,
+                class_name="invoice_date",
+                x_center=0.6,
+                y_center=0.4,
+                width=0.15,
+                height=0.04,
+                bbox_x=200,
+                bbox_y=300,
+                bbox_width=100,
+                bbox_height=25,
+                text_value="2024-01-15",
+                confidence=0.88,
+                source="auto",
+            )
+
+            assert result is not None
+            mock_session.add.assert_called_once()
+            added_annotation = mock_session.add.call_args[0][0]
+            assert added_annotation.text_value == "2024-01-15"
+            assert added_annotation.confidence == 0.88
+            assert added_annotation.source == "auto"
+
+    def test_create_default_source_is_manual(self, repo):
+        """Test create uses manual as default source."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(
+                document_id=str(uuid4()),
+                page_number=1,
+                class_id=0,
+                class_name="invoice_number",
+                x_center=0.5,
+                y_center=0.3,
+                width=0.2,
+                height=0.05,
+                bbox_x=100,
+                bbox_y=200,
+                bbox_width=150,
+                bbox_height=30,
+            )
+
+            added_annotation = mock_session.add.call_args[0][0]
+            assert added_annotation.source == "manual"
+
+    # =========================================================================
+    # create_batch() tests
+    # =========================================================================
+
+    def test_create_batch_returns_ids(self, repo):
+        """Test create_batch returns list of annotation IDs."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            annotations = [
+                {
+                    "document_id": str(uuid4()),
+                    "class_id": 0,
+                    "class_name": "invoice_number",
+                    "x_center": 0.5,
+                    "y_center": 0.3,
+                    "width": 0.2,
+                    "height": 0.05,
+                    "bbox_x": 100,
+                    "bbox_y": 200,
+                    "bbox_width": 150,
+                    "bbox_height": 30,
+                },
+                {
+                    "document_id": str(uuid4()),
+                    "class_id": 1,
+                    "class_name": "invoice_date",
+                    "x_center": 0.6,
+                    "y_center": 0.4,
+                    "width": 0.15,
+                    "height": 0.04,
+                    "bbox_x": 200,
+                    "bbox_y": 300,
+                    "bbox_width": 100,
+                    "bbox_height": 25,
+                },
+            ]
+
+            result = repo.create_batch(annotations)
+
+            assert len(result) == 2
+            assert mock_session.add.call_count == 2
+            assert mock_session.flush.call_count == 2
+
+    def test_create_batch_default_page_number(self, repo):
+        """Test create_batch uses page_number=1 by default."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            annotations = [
+                {
+                    "document_id": str(uuid4()),
+                    "class_id": 0,
+                    "class_name": "invoice_number",
+                    "x_center": 0.5,
+                    "y_center": 0.3,
+                    "width": 0.2,
+                    "height": 0.05,
+                    "bbox_x": 100,
+                    "bbox_y": 200,
+                    "bbox_width": 150,
+                    "bbox_height": 30,
+                    # no page_number
+                },
+            ]
+
+            repo.create_batch(annotations)
+
+            added_annotation = mock_session.add.call_args[0][0]
+            assert added_annotation.page_number == 1
+
+    def test_create_batch_with_all_optional_params(self, repo):
+        """Test create_batch with all optional parameters."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            annotations = [
+                {
+                    "document_id": str(uuid4()),
+                    "page_number": 3,
+                    "class_id": 0,
+                    "class_name": "invoice_number",
+                    "x_center": 0.5,
+                    "y_center": 0.3,
+                    "width": 0.2,
+                    "height": 0.05,
+                    "bbox_x": 100,
+                    "bbox_y": 200,
+                    "bbox_width": 150,
+                    "bbox_height": 30,
+                    "text_value": "INV-123",
+                    "confidence": 0.92,
+                    "source": "ocr",
+                },
+            ]
+
+            repo.create_batch(annotations)
+
+            added_annotation = mock_session.add.call_args[0][0]
+            assert added_annotation.page_number == 3
+            assert added_annotation.text_value == "INV-123"
+            assert added_annotation.confidence == 0.92
+            assert added_annotation.source == "ocr"
+
+    def test_create_batch_empty_list(self, repo):
+        """Test create_batch with empty list returns empty."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create_batch([])
+
+            assert result == []
+            mock_session.add.assert_not_called()
+
+    # =========================================================================
+    # get() tests
+    # =========================================================================
+
+    def test_get_returns_annotation(self, repo, sample_annotation):
+        """Test get returns annotation when exists."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(sample_annotation.annotation_id))
+
+            assert result is not None
+            assert result.class_name == "invoice_number"
+            mock_session.expunge.assert_called_once()
+
+    def test_get_returns_none_when_not_found(self, repo):
+        """Test get returns None when annotation not found."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(uuid4()))
+
+            assert result is None
+            mock_session.expunge.assert_not_called()
+
+    # =========================================================================
+    # get_for_document() tests
+    # =========================================================================
+
+    def test_get_for_document_returns_all_annotations(self, repo, sample_annotation):
+        """Test get_for_document returns all annotations for document."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_annotation]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_for_document(str(sample_annotation.document_id))
+
+            assert len(result) == 1
+            assert result[0].class_name == "invoice_number"
+
+    def test_get_for_document_with_page_filter(self, repo, sample_annotation):
+        """Test get_for_document filters by page number."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_annotation]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_for_document(str(sample_annotation.document_id), page_number=1)
+
+            assert len(result) == 1
+
+    def test_get_for_document_returns_empty_list(self, repo):
+        """Test get_for_document returns empty list when no annotations."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_for_document(str(uuid4()))
+
+            assert result == []
+
+    # =========================================================================
+    # update() tests
+    # =========================================================================
+
+    def test_update_returns_true(self, repo, sample_annotation):
+        """Test update returns True when annotation exists."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(
+                str(sample_annotation.annotation_id),
+                text_value="INV-002",
+            )
+
+            assert result is True
+            assert sample_annotation.text_value == "INV-002"
+
+    def test_update_returns_false_when_not_found(self, repo):
+        """Test update returns False when annotation not found."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(str(uuid4()), text_value="INV-002")
+
+            assert result is False
+
+    def test_update_all_fields(self, repo, sample_annotation):
+        """Test update can update all fields."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(
+                str(sample_annotation.annotation_id),
+                x_center=0.6,
+                y_center=0.4,
+                width=0.25,
+                height=0.06,
+                bbox_x=150,
+                bbox_y=250,
+                bbox_width=175,
+                bbox_height=35,
+                text_value="NEW-VALUE",
+                class_id=5,
+                class_name="new_class",
+            )
+
+            assert result is True
+            assert sample_annotation.x_center == 0.6
+            assert sample_annotation.y_center == 0.4
+            assert sample_annotation.width == 0.25
+            assert sample_annotation.height == 0.06
+            assert sample_annotation.bbox_x == 150
+            assert sample_annotation.bbox_y == 250
+            assert sample_annotation.bbox_width == 175
+            assert sample_annotation.bbox_height == 35
+            assert sample_annotation.text_value == "NEW-VALUE"
+            assert sample_annotation.class_id == 5
+            assert sample_annotation.class_name == "new_class"
+
+    def test_update_partial_fields(self, repo, sample_annotation):
+        """Test update only updates provided fields."""
+        original_x = sample_annotation.x_center
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(
+                str(sample_annotation.annotation_id),
+                text_value="UPDATED",
+            )
+
+            assert result is True
+            assert sample_annotation.text_value == "UPDATED"
+            assert sample_annotation.x_center == original_x  # unchanged
+
+    # =========================================================================
+    # delete() tests
+    # =========================================================================
+
+    def test_delete_returns_true(self, repo, sample_annotation):
+        """Test delete returns True when annotation exists."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(sample_annotation.annotation_id))
+
+            assert result is True
+            mock_session.delete.assert_called_once()
+
+    def test_delete_returns_false_when_not_found(self, repo):
+        """Test delete returns False when annotation not found."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(uuid4()))
+
+            assert result is False
+            mock_session.delete.assert_not_called()
+
+    # =========================================================================
+    # delete_for_document() tests
+    # =========================================================================
+
+    def test_delete_for_document_returns_count(self, repo, sample_annotation):
+        """Test delete_for_document returns count of deleted annotations."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_annotation]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete_for_document(str(sample_annotation.document_id))
+
+            assert result == 1
+            mock_session.delete.assert_called_once()
+
+    def test_delete_for_document_with_source_filter(self, repo, sample_annotation):
+        """Test delete_for_document filters by source."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_annotation]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete_for_document(str(sample_annotation.document_id), source="auto")
+
+            assert result == 1
+
+    def test_delete_for_document_returns_zero(self, repo):
+        """Test delete_for_document returns 0 when no annotations."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete_for_document(str(uuid4()))
+
+            assert result == 0
+            mock_session.delete.assert_not_called()
+
+    # =========================================================================
+    # verify() tests
+    # =========================================================================
+
+    def test_verify_marks_annotation_verified(self, repo, sample_annotation):
+        """Test verify marks annotation as verified."""
+        sample_annotation.is_verified = False
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.verify(str(sample_annotation.annotation_id), "admin-token")
+
+            assert result is not None
+            assert sample_annotation.is_verified is True
+            assert sample_annotation.verified_by == "admin-token"
+            mock_session.commit.assert_called_once()
+
+    def test_verify_returns_none_when_not_found(self, repo):
+        """Test verify returns None when annotation not found."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.verify(str(uuid4()), "admin-token")
+
+            assert result is None
+
+    # =========================================================================
+    # override() tests
+    # =========================================================================
+
+    def test_override_updates_annotation(self, repo, sample_annotation):
+        """Test override updates annotation and creates history."""
+        sample_annotation.source = "auto"
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.override(
+                str(sample_annotation.annotation_id),
+                "admin-token",
+                change_reason="Correction",
+                text_value="NEW-VALUE",
+            )
+
+            assert result is not None
+            assert sample_annotation.text_value == "NEW-VALUE"
+            assert sample_annotation.source == "manual"
+            assert sample_annotation.override_source == "auto"
+            assert mock_session.add.call_count >= 2  # annotation + history
+
+    def test_override_returns_none_when_not_found(self, repo):
+        """Test override returns None when annotation not found."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.override(str(uuid4()), "admin-token", text_value="NEW")
+
+            assert result is None
+
+    def test_override_does_not_change_source_if_already_manual(self, repo, sample_annotation):
+        """Test override does not change override_source if already manual."""
+        sample_annotation.source = "manual"
+        sample_annotation.override_source = None
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.override(
+                str(sample_annotation.annotation_id),
+                "admin-token",
+                text_value="NEW-VALUE",
+            )
+
+            assert sample_annotation.source == "manual"
+            assert sample_annotation.override_source is None
+
+    def test_override_skips_unknown_attributes(self, repo, sample_annotation):
+        """Test override ignores unknown attributes."""
+        sample_annotation.source = "auto"
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_annotation
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.override(
+                str(sample_annotation.annotation_id),
+                "admin-token",
+                unknown_field="should_be_ignored",
+                text_value="VALID",
+            )
+
+            assert result is not None
+            assert sample_annotation.text_value == "VALID"
+            assert not hasattr(sample_annotation, "unknown_field") or getattr(sample_annotation, "unknown_field", None) != "should_be_ignored"
+
+    # =========================================================================
+    # create_history() tests
+    # =========================================================================
+
+    def test_create_history_returns_history(self, repo):
+        """Test create_history returns created history record."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            annotation_id = uuid4()
+            document_id = uuid4()
+            result = repo.create_history(
+                annotation_id=annotation_id,
+                document_id=document_id,
+                action="create",
+                previous_value=None,
+                new_value={"class_name": "invoice_number"},
+                changed_by="admin-token",
+                change_reason="Initial creation",
+            )
+
+            mock_session.add.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_create_history_with_minimal_params(self, repo):
+        """Test create_history with minimal parameters."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create_history(
+                annotation_id=uuid4(),
+                document_id=uuid4(),
+                action="delete",
+            )
+
+            mock_session.add.assert_called_once()
+            added_history = mock_session.add.call_args[0][0]
+            assert added_history.action == "delete"
+            assert added_history.previous_value is None
+            assert added_history.new_value is None
+
+    # =========================================================================
+    # get_history() tests
+    # =========================================================================
+
+    def test_get_history_returns_list(self, repo, sample_history):
+        """Test get_history returns list of history records."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_history]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_history(sample_history.annotation_id)
+
+            assert len(result) == 1
+            assert result[0].action == "override"
+
+    def test_get_history_returns_empty_list(self, repo):
+        """Test get_history returns empty list when no history."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_history(uuid4())
+
+            assert result == []
+
+    # =========================================================================
+    # get_document_history() tests
+    # =========================================================================
+
+    def test_get_document_history_returns_list(self, repo, sample_history):
+        """Test get_document_history returns list of history records."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_history]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_document_history(sample_history.document_id)
+
+            assert len(result) == 1
+
+    def test_get_document_history_returns_empty_list(self, repo):
+        """Test get_document_history returns empty list when no history."""
+        with patch("inference.data.repositories.annotation_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_document_history(uuid4())
+
+            assert result == []

tests/data/repositories/test_base_repository.py

@@ -0,0 +1,142 @@
+"""
+Tests for BaseRepository
+
+100% coverage tests for base repository utilities.
+"""
+
+import pytest
+from datetime import datetime, timezone
+from unittest.mock import MagicMock, patch
+from uuid import uuid4, UUID
+
+from inference.data.repositories.base import BaseRepository
+
+
+class ConcreteRepository(BaseRepository[MagicMock]):
+    """Concrete implementation for testing abstract base class."""
+    pass
+
+
+class TestBaseRepository:
+    """Tests for BaseRepository."""
+
+    @pytest.fixture
+    def repo(self) -> ConcreteRepository:
+        """Create a ConcreteRepository instance."""
+        return ConcreteRepository()
+
+    # =========================================================================
+    # _session() tests
+    # =========================================================================
+
+    def test_session_yields_session(self, repo):
+        """Test _session yields a database session."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            with repo._session() as session:
+                assert session is mock_session
+
+    # =========================================================================
+    # _expunge() tests
+    # =========================================================================
+
+    def test_expunge_detaches_entity(self, repo):
+        """Test _expunge detaches entity from session."""
+        mock_session = MagicMock()
+        mock_entity = MagicMock()
+
+        result = repo._expunge(mock_session, mock_entity)
+
+        mock_session.expunge.assert_called_once_with(mock_entity)
+        assert result is mock_entity
+
+    # =========================================================================
+    # _expunge_all() tests
+    # =========================================================================
+
+    def test_expunge_all_detaches_all_entities(self, repo):
+        """Test _expunge_all detaches all entities from session."""
+        mock_session = MagicMock()
+        mock_entity1 = MagicMock()
+        mock_entity2 = MagicMock()
+        entities = [mock_entity1, mock_entity2]
+
+        result = repo._expunge_all(mock_session, entities)
+
+        assert mock_session.expunge.call_count == 2
+        mock_session.expunge.assert_any_call(mock_entity1)
+        mock_session.expunge.assert_any_call(mock_entity2)
+        assert result is entities
+
+    def test_expunge_all_empty_list(self, repo):
+        """Test _expunge_all with empty list."""
+        mock_session = MagicMock()
+        entities = []
+
+        result = repo._expunge_all(mock_session, entities)
+
+        mock_session.expunge.assert_not_called()
+        assert result == []
+
+    # =========================================================================
+    # _now() tests
+    # =========================================================================
+
+    def test_now_returns_utc_datetime(self, repo):
+        """Test _now returns timezone-aware UTC datetime."""
+        result = repo._now()
+
+        assert result.tzinfo == timezone.utc
+        assert isinstance(result, datetime)
+
+    def test_now_is_recent(self, repo):
+        """Test _now returns a recent datetime."""
+        before = datetime.now(timezone.utc)
+        result = repo._now()
+        after = datetime.now(timezone.utc)
+
+        assert before <= result <= after
+
+    # =========================================================================
+    # _validate_uuid() tests
+    # =========================================================================
+
+    def test_validate_uuid_with_valid_string(self, repo):
+        """Test _validate_uuid with valid UUID string."""
+        valid_uuid_str = str(uuid4())
+
+        result = repo._validate_uuid(valid_uuid_str)
+
+        assert isinstance(result, UUID)
+        assert str(result) == valid_uuid_str
+
+    def test_validate_uuid_with_invalid_string(self, repo):
+        """Test _validate_uuid raises ValueError for invalid UUID."""
+        with pytest.raises(ValueError) as exc_info:
+            repo._validate_uuid("not-a-valid-uuid")
+
+        assert "Invalid id" in str(exc_info.value)
+
+    def test_validate_uuid_with_custom_field_name(self, repo):
+        """Test _validate_uuid uses custom field name in error."""
+        with pytest.raises(ValueError) as exc_info:
+            repo._validate_uuid("invalid", field_name="document_id")
+
+        assert "Invalid document_id" in str(exc_info.value)
+
+    def test_validate_uuid_with_none(self, repo):
+        """Test _validate_uuid raises ValueError for None."""
+        with pytest.raises(ValueError) as exc_info:
+            repo._validate_uuid(None)
+
+        assert "Invalid id" in str(exc_info.value)
+
+    def test_validate_uuid_with_empty_string(self, repo):
+        """Test _validate_uuid raises ValueError for empty string."""
+        with pytest.raises(ValueError) as exc_info:
+            repo._validate_uuid("")
+
+        assert "Invalid id" in str(exc_info.value)

tests/data/repositories/test_batch_upload_repository.py

@@ -0,0 +1,386 @@
+"""
+Tests for BatchUploadRepository
+
+100% coverage tests for batch upload management.
+"""
+
+import pytest
+from datetime import datetime, timezone
+from unittest.mock import MagicMock, patch
+from uuid import uuid4, UUID
+
+from inference.data.admin_models import BatchUpload, BatchUploadFile
+from inference.data.repositories.batch_upload_repository import BatchUploadRepository
+
+
+class TestBatchUploadRepository:
+    """Tests for BatchUploadRepository."""
+
+    @pytest.fixture
+    def sample_batch(self) -> BatchUpload:
+        """Create a sample batch upload for testing."""
+        return BatchUpload(
+            batch_id=uuid4(),
+            admin_token="admin-token",
+            filename="invoices.zip",
+            file_size=1024000,
+            upload_source="ui",
+            status="pending",
+            total_files=10,
+            processed_files=0,
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def sample_file(self) -> BatchUploadFile:
+        """Create a sample batch upload file for testing."""
+        return BatchUploadFile(
+            file_id=uuid4(),
+            batch_id=uuid4(),
+            filename="invoice_001.pdf",
+            status="pending",
+            created_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def repo(self) -> BatchUploadRepository:
+        """Create a BatchUploadRepository instance."""
+        return BatchUploadRepository()
+
+    # =========================================================================
+    # create() tests
+    # =========================================================================
+
+    def test_create_returns_batch(self, repo):
+        """Test create returns created batch upload."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                admin_token="admin-token",
+                filename="test.zip",
+                file_size=1024,
+            )
+
+            mock_session.add.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_create_with_upload_source(self, repo):
+        """Test create with custom upload source."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(
+                admin_token="admin-token",
+                filename="test.zip",
+                file_size=1024,
+                upload_source="api",
+            )
+
+            added_batch = mock_session.add.call_args[0][0]
+            assert added_batch.upload_source == "api"
+
+    def test_create_default_upload_source(self, repo):
+        """Test create uses default upload source."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(
+                admin_token="admin-token",
+                filename="test.zip",
+                file_size=1024,
+            )
+
+            added_batch = mock_session.add.call_args[0][0]
+            assert added_batch.upload_source == "ui"
+
+    # =========================================================================
+    # get() tests
+    # =========================================================================
+
+    def test_get_returns_batch(self, repo, sample_batch):
+        """Test get returns batch when exists."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_batch
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(sample_batch.batch_id)
+
+            assert result is not None
+            assert result.filename == "invoices.zip"
+            mock_session.expunge.assert_called_once()
+
+    def test_get_returns_none_when_not_found(self, repo):
+        """Test get returns None when batch not found."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(uuid4())
+
+            assert result is None
+            mock_session.expunge.assert_not_called()
+
+    # =========================================================================
+    # update() tests
+    # =========================================================================
+
+    def test_update_updates_batch(self, repo, sample_batch):
+        """Test update updates batch fields."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_batch
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update(
+                sample_batch.batch_id,
+                status="processing",
+                processed_files=5,
+            )
+
+            assert sample_batch.status == "processing"
+            assert sample_batch.processed_files == 5
+            mock_session.add.assert_called_once()
+
+    def test_update_ignores_unknown_fields(self, repo, sample_batch):
+        """Test update ignores unknown fields."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_batch
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update(
+                sample_batch.batch_id,
+                unknown_field="should_be_ignored",
+            )
+
+            mock_session.add.assert_called_once()
+
+    def test_update_not_found(self, repo):
+        """Test update does nothing when batch not found."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update(uuid4(), status="processing")
+
+            mock_session.add.assert_not_called()
+
+    def test_update_multiple_fields(self, repo, sample_batch):
+        """Test update can update multiple fields."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_batch
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update(
+                sample_batch.batch_id,
+                status="completed",
+                processed_files=10,
+                total_files=10,
+            )
+
+            assert sample_batch.status == "completed"
+            assert sample_batch.processed_files == 10
+            assert sample_batch.total_files == 10
+
+    # =========================================================================
+    # create_file() tests
+    # =========================================================================
+
+    def test_create_file_returns_file(self, repo):
+        """Test create_file returns created file record."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create_file(
+                batch_id=uuid4(),
+                filename="invoice_001.pdf",
+            )
+
+            mock_session.add.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_create_file_with_kwargs(self, repo):
+        """Test create_file with additional kwargs."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create_file(
+                batch_id=uuid4(),
+                filename="invoice_001.pdf",
+                status="processing",
+                file_size=1024,
+            )
+
+            added_file = mock_session.add.call_args[0][0]
+            assert added_file.filename == "invoice_001.pdf"
+
+    # =========================================================================
+    # update_file() tests
+    # =========================================================================
+
+    def test_update_file_updates_file(self, repo, sample_file):
+        """Test update_file updates file fields."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_file
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_file(
+                sample_file.file_id,
+                status="completed",
+            )
+
+            assert sample_file.status == "completed"
+            mock_session.add.assert_called_once()
+
+    def test_update_file_ignores_unknown_fields(self, repo, sample_file):
+        """Test update_file ignores unknown fields."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_file
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_file(
+                sample_file.file_id,
+                unknown_field="should_be_ignored",
+            )
+
+            mock_session.add.assert_called_once()
+
+    def test_update_file_not_found(self, repo):
+        """Test update_file does nothing when file not found."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_file(uuid4(), status="completed")
+
+            mock_session.add.assert_not_called()
+
+    def test_update_file_multiple_fields(self, repo, sample_file):
+        """Test update_file can update multiple fields."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_file
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_file(
+                sample_file.file_id,
+                status="failed",
+            )
+
+            assert sample_file.status == "failed"
+
+    # =========================================================================
+    # get_files() tests
+    # =========================================================================
+
+    def test_get_files_returns_list(self, repo, sample_file):
+        """Test get_files returns list of files."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_file]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_files(sample_file.batch_id)
+
+            assert len(result) == 1
+            assert result[0].filename == "invoice_001.pdf"
+
+    def test_get_files_returns_empty_list(self, repo):
+        """Test get_files returns empty list when no files."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_files(uuid4())
+
+            assert result == []
+
+    # =========================================================================
+    # get_paginated() tests
+    # =========================================================================
+
+    def test_get_paginated_returns_batches_and_total(self, repo, sample_batch):
+        """Test get_paginated returns list of batches and total count."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_batch]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            batches, total = repo.get_paginated()
+
+            assert len(batches) == 1
+            assert total == 1
+
+    def test_get_paginated_with_pagination(self, repo, sample_batch):
+        """Test get_paginated with limit and offset."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 100
+            mock_session.exec.return_value.all.return_value = [sample_batch]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            batches, total = repo.get_paginated(limit=25, offset=50)
+
+            assert total == 100
+
+    def test_get_paginated_empty_results(self, repo):
+        """Test get_paginated with no results."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 0
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            batches, total = repo.get_paginated()
+
+            assert batches == []
+            assert total == 0
+
+    def test_get_paginated_with_admin_token(self, repo, sample_batch):
+        """Test get_paginated with admin_token parameter (deprecated, ignored)."""
+        with patch("inference.data.repositories.batch_upload_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_batch]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            batches, total = repo.get_paginated(admin_token="admin-token")
+
+            assert len(batches) == 1

tests/data/repositories/test_dataset_repository.py

@@ -0,0 +1,597 @@
+"""
+Tests for DatasetRepository
+
+100% coverage tests for dataset management.
+"""
+
+import pytest
+from datetime import datetime, timezone
+from unittest.mock import MagicMock, patch
+from uuid import uuid4, UUID
+
+from inference.data.admin_models import TrainingDataset, DatasetDocument, TrainingTask
+from inference.data.repositories.dataset_repository import DatasetRepository
+
+
+class TestDatasetRepository:
+    """Tests for DatasetRepository."""
+
+    @pytest.fixture
+    def sample_dataset(self) -> TrainingDataset:
+        """Create a sample dataset for testing."""
+        return TrainingDataset(
+            dataset_id=uuid4(),
+            name="Test Dataset",
+            description="A test dataset",
+            status="ready",
+            train_ratio=0.8,
+            val_ratio=0.1,
+            seed=42,
+            total_documents=100,
+            total_images=100,
+            total_annotations=500,
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def sample_dataset_document(self) -> DatasetDocument:
+        """Create a sample dataset document for testing."""
+        return DatasetDocument(
+            id=uuid4(),
+            dataset_id=uuid4(),
+            document_id=uuid4(),
+            split="train",
+            page_count=2,
+            annotation_count=10,
+            created_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def sample_training_task(self) -> TrainingTask:
+        """Create a sample training task for testing."""
+        return TrainingTask(
+            task_id=uuid4(),
+            admin_token="admin-token",
+            name="Test Task",
+            status="running",
+            dataset_id=uuid4(),
+        )
+
+    @pytest.fixture
+    def repo(self) -> DatasetRepository:
+        """Create a DatasetRepository instance."""
+        return DatasetRepository()
+
+    # =========================================================================
+    # create() tests
+    # =========================================================================
+
+    def test_create_returns_dataset(self, repo):
+        """Test create returns created dataset."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(name="Test Dataset")
+
+            mock_session.add.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_create_with_all_params(self, repo):
+        """Test create with all parameters."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                name="Full Dataset",
+                description="A complete dataset",
+                train_ratio=0.7,
+                val_ratio=0.15,
+                seed=123,
+            )
+
+            added_dataset = mock_session.add.call_args[0][0]
+            assert added_dataset.name == "Full Dataset"
+            assert added_dataset.description == "A complete dataset"
+            assert added_dataset.train_ratio == 0.7
+            assert added_dataset.val_ratio == 0.15
+            assert added_dataset.seed == 123
+
+    def test_create_default_values(self, repo):
+        """Test create uses default values."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(name="Minimal Dataset")
+
+            added_dataset = mock_session.add.call_args[0][0]
+            assert added_dataset.train_ratio == 0.8
+            assert added_dataset.val_ratio == 0.1
+            assert added_dataset.seed == 42
+
+    # =========================================================================
+    # get() tests
+    # =========================================================================
+
+    def test_get_returns_dataset(self, repo, sample_dataset):
+        """Test get returns dataset when exists."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(sample_dataset.dataset_id))
+
+            assert result is not None
+            assert result.name == "Test Dataset"
+            mock_session.expunge.assert_called_once()
+
+    def test_get_with_uuid(self, repo, sample_dataset):
+        """Test get works with UUID object."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(sample_dataset.dataset_id)
+
+            assert result is not None
+
+    def test_get_returns_none_when_not_found(self, repo):
+        """Test get returns None when dataset not found."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(uuid4()))
+
+            assert result is None
+            mock_session.expunge.assert_not_called()
+
+    # =========================================================================
+    # get_paginated() tests
+    # =========================================================================
+
+    def test_get_paginated_returns_datasets_and_total(self, repo, sample_dataset):
+        """Test get_paginated returns list of datasets and total count."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_dataset]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            datasets, total = repo.get_paginated()
+
+            assert len(datasets) == 1
+            assert total == 1
+
+    def test_get_paginated_with_status_filter(self, repo, sample_dataset):
+        """Test get_paginated filters by status."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_dataset]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            datasets, total = repo.get_paginated(status="ready")
+
+            assert len(datasets) == 1
+
+    def test_get_paginated_with_pagination(self, repo, sample_dataset):
+        """Test get_paginated with limit and offset."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 50
+            mock_session.exec.return_value.all.return_value = [sample_dataset]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            datasets, total = repo.get_paginated(limit=10, offset=20)
+
+            assert total == 50
+
+    def test_get_paginated_empty_results(self, repo):
+        """Test get_paginated with no results."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 0
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            datasets, total = repo.get_paginated()
+
+            assert datasets == []
+            assert total == 0
+
+    # =========================================================================
+    # get_active_training_tasks() tests
+    # =========================================================================
+
+    def test_get_active_training_tasks_returns_dict(self, repo, sample_training_task):
+        """Test get_active_training_tasks returns dict of active tasks."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_training_task]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_active_training_tasks([str(sample_training_task.dataset_id)])
+
+            assert str(sample_training_task.dataset_id) in result
+
+    def test_get_active_training_tasks_empty_input(self, repo):
+        """Test get_active_training_tasks with empty input."""
+        result = repo.get_active_training_tasks([])
+
+        assert result == {}
+
+    def test_get_active_training_tasks_invalid_uuid(self, repo):
+        """Test get_active_training_tasks filters invalid UUIDs."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_active_training_tasks(["invalid-uuid", str(uuid4())])
+
+            # Should still query with valid UUID
+            assert result == {}
+
+    def test_get_active_training_tasks_all_invalid_uuids(self, repo):
+        """Test get_active_training_tasks with all invalid UUIDs."""
+        result = repo.get_active_training_tasks(["invalid-uuid-1", "invalid-uuid-2"])
+
+        assert result == {}
+
+    # =========================================================================
+    # update_status() tests
+    # =========================================================================
+
+    def test_update_status_updates_dataset(self, repo, sample_dataset):
+        """Test update_status updates dataset status."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(sample_dataset.dataset_id), "training")
+
+            assert sample_dataset.status == "training"
+            mock_session.commit.assert_called_once()
+
+    def test_update_status_with_error_message(self, repo, sample_dataset):
+        """Test update_status with error message."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(
+                str(sample_dataset.dataset_id),
+                "failed",
+                error_message="Training failed",
+            )
+
+            assert sample_dataset.error_message == "Training failed"
+
+    def test_update_status_with_totals(self, repo, sample_dataset):
+        """Test update_status with total counts."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(
+                str(sample_dataset.dataset_id),
+                "ready",
+                total_documents=200,
+                total_images=200,
+                total_annotations=1000,
+            )
+
+            assert sample_dataset.total_documents == 200
+            assert sample_dataset.total_images == 200
+            assert sample_dataset.total_annotations == 1000
+
+    def test_update_status_with_dataset_path(self, repo, sample_dataset):
+        """Test update_status with dataset path."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(
+                str(sample_dataset.dataset_id),
+                "ready",
+                dataset_path="/path/to/dataset",
+            )
+
+            assert sample_dataset.dataset_path == "/path/to/dataset"
+
+    def test_update_status_with_uuid(self, repo, sample_dataset):
+        """Test update_status works with UUID object."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(sample_dataset.dataset_id, "ready")
+
+            assert sample_dataset.status == "ready"
+
+    def test_update_status_not_found(self, repo):
+        """Test update_status does nothing when dataset not found."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(uuid4()), "ready")
+
+            mock_session.add.assert_not_called()
+
+    # =========================================================================
+    # update_training_status() tests
+    # =========================================================================
+
+    def test_update_training_status_updates_dataset(self, repo, sample_dataset):
+        """Test update_training_status updates training status."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_training_status(str(sample_dataset.dataset_id), "running")
+
+            assert sample_dataset.training_status == "running"
+            mock_session.commit.assert_called_once()
+
+    def test_update_training_status_with_task_id(self, repo, sample_dataset):
+        """Test update_training_status with active task ID."""
+        task_id = uuid4()
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_training_status(
+                str(sample_dataset.dataset_id),
+                "running",
+                active_training_task_id=str(task_id),
+            )
+
+            assert sample_dataset.active_training_task_id == task_id
+
+    def test_update_training_status_updates_main_status(self, repo, sample_dataset):
+        """Test update_training_status updates main status when completed."""
+        sample_dataset.status = "ready"
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_training_status(
+                str(sample_dataset.dataset_id),
+                "completed",
+                update_main_status=True,
+            )
+
+            assert sample_dataset.training_status == "completed"
+            assert sample_dataset.status == "trained"
+
+    def test_update_training_status_clears_task_id(self, repo, sample_dataset):
+        """Test update_training_status clears task ID when None."""
+        sample_dataset.active_training_task_id = uuid4()
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_training_status(
+                str(sample_dataset.dataset_id),
+                None,
+                active_training_task_id=None,
+            )
+
+            assert sample_dataset.active_training_task_id is None
+
+    def test_update_training_status_not_found(self, repo):
+        """Test update_training_status does nothing when dataset not found."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_training_status(str(uuid4()), "running")
+
+            mock_session.add.assert_not_called()
+
+    # =========================================================================
+    # add_documents() tests
+    # =========================================================================
+
+    def test_add_documents_creates_links(self, repo):
+        """Test add_documents creates dataset document links."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            documents = [
+                {
+                    "document_id": str(uuid4()),
+                    "split": "train",
+                    "page_count": 2,
+                    "annotation_count": 10,
+                },
+                {
+                    "document_id": str(uuid4()),
+                    "split": "val",
+                    "page_count": 1,
+                    "annotation_count": 5,
+                },
+            ]
+
+            repo.add_documents(str(uuid4()), documents)
+
+            assert mock_session.add.call_count == 2
+            mock_session.commit.assert_called_once()
+
+    def test_add_documents_default_counts(self, repo):
+        """Test add_documents uses default counts."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            documents = [
+                {
+                    "document_id": str(uuid4()),
+                    "split": "train",
+                },
+            ]
+
+            repo.add_documents(str(uuid4()), documents)
+
+            added_doc = mock_session.add.call_args[0][0]
+            assert added_doc.page_count == 0
+            assert added_doc.annotation_count == 0
+
+    def test_add_documents_with_uuid(self, repo):
+        """Test add_documents works with UUID object."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            documents = [
+                {
+                    "document_id": uuid4(),
+                    "split": "train",
+                },
+            ]
+
+            repo.add_documents(uuid4(), documents)
+
+            mock_session.add.assert_called_once()
+
+    def test_add_documents_empty_list(self, repo):
+        """Test add_documents with empty list."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.add_documents(str(uuid4()), [])
+
+            mock_session.add.assert_not_called()
+            mock_session.commit.assert_called_once()
+
+    # =========================================================================
+    # get_documents() tests
+    # =========================================================================
+
+    def test_get_documents_returns_list(self, repo, sample_dataset_document):
+        """Test get_documents returns list of dataset documents."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_dataset_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_documents(str(sample_dataset_document.dataset_id))
+
+            assert len(result) == 1
+            assert result[0].split == "train"
+
+    def test_get_documents_with_uuid(self, repo, sample_dataset_document):
+        """Test get_documents works with UUID object."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_dataset_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_documents(sample_dataset_document.dataset_id)
+
+            assert len(result) == 1
+
+    def test_get_documents_returns_empty_list(self, repo):
+        """Test get_documents returns empty list when no documents."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_documents(str(uuid4()))
+
+            assert result == []
+
+    # =========================================================================
+    # delete() tests
+    # =========================================================================
+
+    def test_delete_returns_true(self, repo, sample_dataset):
+        """Test delete returns True when dataset exists."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(sample_dataset.dataset_id))
+
+            assert result is True
+            mock_session.delete.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_delete_with_uuid(self, repo, sample_dataset):
+        """Test delete works with UUID object."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_dataset
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(sample_dataset.dataset_id)
+
+            assert result is True
+
+    def test_delete_returns_false_when_not_found(self, repo):
+        """Test delete returns False when dataset not found."""
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(uuid4()))
+
+            assert result is False
+            mock_session.delete.assert_not_called()

tests/data/repositories/test_document_repository.py

@@ -0,0 +1,748 @@
+"""
+Tests for DocumentRepository
+
+Comprehensive TDD tests for document management - targeting 100% coverage.
+"""
+
+import pytest
+from datetime import datetime, timedelta, timezone
+from unittest.mock import MagicMock, patch
+from uuid import uuid4
+
+from inference.data.admin_models import AdminDocument, AdminAnnotation
+from inference.data.repositories.document_repository import DocumentRepository
+
+
+class TestDocumentRepository:
+    """Tests for DocumentRepository."""
+
+    @pytest.fixture
+    def sample_document(self) -> AdminDocument:
+        """Create a sample document for testing."""
+        return AdminDocument(
+            document_id=uuid4(),
+            filename="test.pdf",
+            file_size=1024,
+            content_type="application/pdf",
+            file_path="/tmp/test.pdf",
+            page_count=1,
+            status="pending",
+            category="invoice",
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def labeled_document(self) -> AdminDocument:
+        """Create a labeled document for testing."""
+        return AdminDocument(
+            document_id=uuid4(),
+            filename="labeled.pdf",
+            file_size=2048,
+            content_type="application/pdf",
+            file_path="/tmp/labeled.pdf",
+            page_count=2,
+            status="labeled",
+            category="invoice",
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def locked_document(self) -> AdminDocument:
+        """Create a document with annotation lock."""
+        doc = AdminDocument(
+            document_id=uuid4(),
+            filename="locked.pdf",
+            file_size=1024,
+            content_type="application/pdf",
+            file_path="/tmp/locked.pdf",
+            page_count=1,
+            status="pending",
+            category="invoice",
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+        doc.annotation_lock_until = datetime.now(timezone.utc) + timedelta(minutes=5)
+        return doc
+
+    @pytest.fixture
+    def expired_lock_document(self) -> AdminDocument:
+        """Create a document with expired annotation lock."""
+        doc = AdminDocument(
+            document_id=uuid4(),
+            filename="expired_lock.pdf",
+            file_size=1024,
+            content_type="application/pdf",
+            file_path="/tmp/expired_lock.pdf",
+            page_count=1,
+            status="pending",
+            category="invoice",
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+        doc.annotation_lock_until = datetime.now(timezone.utc) - timedelta(minutes=5)
+        return doc
+
+    @pytest.fixture
+    def repo(self) -> DocumentRepository:
+        """Create a DocumentRepository instance."""
+        return DocumentRepository()
+
+    # ==========================================================================
+    # create() tests
+    # ==========================================================================
+
+    def test_create_returns_document_id(self, repo):
+        """Test create returns document ID."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                filename="test.pdf",
+                file_size=1024,
+                content_type="application/pdf",
+                file_path="/tmp/test.pdf",
+            )
+
+            assert result is not None
+            mock_session.add.assert_called_once()
+            mock_session.flush.assert_called_once()
+
+    def test_create_with_all_parameters(self, repo):
+        """Test create with all optional parameters."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                filename="test.pdf",
+                file_size=1024,
+                content_type="application/pdf",
+                file_path="/tmp/test.pdf",
+                page_count=5,
+                upload_source="api",
+                csv_field_values={"InvoiceNumber": "INV-001"},
+                group_key="batch-001",
+                category="receipt",
+                admin_token="token-123",
+            )
+
+            assert result is not None
+            added_doc = mock_session.add.call_args[0][0]
+            assert added_doc.page_count == 5
+            assert added_doc.upload_source == "api"
+            assert added_doc.csv_field_values == {"InvoiceNumber": "INV-001"}
+            assert added_doc.group_key == "batch-001"
+            assert added_doc.category == "receipt"
+
+    # ==========================================================================
+    # get() tests
+    # ==========================================================================
+
+    def test_get_returns_document(self, repo, sample_document):
+        """Test get returns document when exists."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(sample_document.document_id))
+
+            assert result is not None
+            assert result.filename == "test.pdf"
+            mock_session.expunge.assert_called_once()
+
+    def test_get_returns_none_when_not_found(self, repo):
+        """Test get returns None when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(uuid4()))
+
+            assert result is None
+
+    # ==========================================================================
+    # get_by_token() tests
+    # ==========================================================================
+
+    def test_get_by_token_delegates_to_get(self, repo, sample_document):
+        """Test get_by_token delegates to get method."""
+        with patch.object(repo, "get", return_value=sample_document) as mock_get:
+            result = repo.get_by_token(str(sample_document.document_id), "token-123")
+
+            assert result == sample_document
+            mock_get.assert_called_once_with(str(sample_document.document_id))
+
+    # ==========================================================================
+    # get_paginated() tests
+    # ==========================================================================
+
+    def test_get_paginated_no_filters(self, repo, sample_document):
+        """Test get_paginated with no filters."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated()
+
+            assert total == 1
+            assert len(results) == 1
+
+    def test_get_paginated_with_status_filter(self, repo, sample_document):
+        """Test get_paginated with status filter."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated(status="pending")
+
+            assert total == 1
+
+    def test_get_paginated_with_upload_source_filter(self, repo, sample_document):
+        """Test get_paginated with upload_source filter."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated(upload_source="ui")
+
+            assert total == 1
+
+    def test_get_paginated_with_auto_label_status_filter(self, repo, sample_document):
+        """Test get_paginated with auto_label_status filter."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated(auto_label_status="completed")
+
+            assert total == 1
+
+    def test_get_paginated_with_batch_id_filter(self, repo, sample_document):
+        """Test get_paginated with batch_id filter."""
+        batch_id = str(uuid4())
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated(batch_id=batch_id)
+
+            assert total == 1
+
+    def test_get_paginated_with_category_filter(self, repo, sample_document):
+        """Test get_paginated with category filter."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated(category="invoice")
+
+            assert total == 1
+
+    def test_get_paginated_with_has_annotations_true(self, repo, sample_document):
+        """Test get_paginated with has_annotations=True."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated(has_annotations=True)
+
+            assert total == 1
+
+    def test_get_paginated_with_has_annotations_false(self, repo, sample_document):
+        """Test get_paginated with has_annotations=False."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_paginated(has_annotations=False)
+
+            assert total == 1
+
+    # ==========================================================================
+    # update_status() tests
+    # ==========================================================================
+
+    def test_update_status(self, repo, sample_document):
+        """Test update_status updates document status."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(sample_document.document_id), "labeled")
+
+            assert sample_document.status == "labeled"
+            mock_session.add.assert_called_once()
+
+    def test_update_status_with_auto_label_status(self, repo, sample_document):
+        """Test update_status with auto_label_status."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(
+                str(sample_document.document_id),
+                "labeled",
+                auto_label_status="completed",
+            )
+
+            assert sample_document.auto_label_status == "completed"
+
+    def test_update_status_with_auto_label_error(self, repo, sample_document):
+        """Test update_status with auto_label_error."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(
+                str(sample_document.document_id),
+                "failed",
+                auto_label_error="OCR failed",
+            )
+
+            assert sample_document.auto_label_error == "OCR failed"
+
+    def test_update_status_document_not_found(self, repo):
+        """Test update_status when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(uuid4()), "labeled")
+
+            mock_session.add.assert_not_called()
+
+    # ==========================================================================
+    # update_file_path() tests
+    # ==========================================================================
+
+    def test_update_file_path(self, repo, sample_document):
+        """Test update_file_path updates document file path."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_file_path(str(sample_document.document_id), "/new/path.pdf")
+
+            assert sample_document.file_path == "/new/path.pdf"
+            mock_session.add.assert_called_once()
+
+    def test_update_file_path_document_not_found(self, repo):
+        """Test update_file_path when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_file_path(str(uuid4()), "/new/path.pdf")
+
+            mock_session.add.assert_not_called()
+
+    # ==========================================================================
+    # update_group_key() tests
+    # ==========================================================================
+
+    def test_update_group_key_returns_true(self, repo, sample_document):
+        """Test update_group_key returns True when document exists."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update_group_key(str(sample_document.document_id), "new-group")
+
+            assert result is True
+            assert sample_document.group_key == "new-group"
+
+    def test_update_group_key_returns_false(self, repo):
+        """Test update_group_key returns False when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update_group_key(str(uuid4()), "new-group")
+
+            assert result is False
+
+    # ==========================================================================
+    # update_category() tests
+    # ==========================================================================
+
+    def test_update_category(self, repo, sample_document):
+        """Test update_category updates document category."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update_category(str(sample_document.document_id), "receipt")
+
+            assert sample_document.category == "receipt"
+            mock_session.add.assert_called()
+
+    def test_update_category_returns_none_when_not_found(self, repo):
+        """Test update_category returns None when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update_category(str(uuid4()), "receipt")
+
+            assert result is None
+
+    # ==========================================================================
+    # delete() tests
+    # ==========================================================================
+
+    def test_delete_returns_true_when_exists(self, repo, sample_document):
+        """Test delete returns True when document exists."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(sample_document.document_id))
+
+            assert result is True
+            mock_session.delete.assert_called_once_with(sample_document)
+
+    def test_delete_with_annotations(self, repo, sample_document):
+        """Test delete removes annotations before deleting document."""
+        annotation = MagicMock()
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_session.exec.return_value.all.return_value = [annotation]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(sample_document.document_id))
+
+            assert result is True
+            assert mock_session.delete.call_count == 2
+
+    def test_delete_returns_false_when_not_exists(self, repo):
+        """Test delete returns False when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(uuid4()))
+
+            assert result is False
+
+    # ==========================================================================
+    # get_categories() tests
+    # ==========================================================================
+
+    def test_get_categories(self, repo):
+        """Test get_categories returns unique categories."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = ["invoice", "receipt", None]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_categories()
+
+            assert result == ["invoice", "receipt"]
+
+    # ==========================================================================
+    # get_labeled_for_export() tests
+    # ==========================================================================
+
+    def test_get_labeled_for_export(self, repo, labeled_document):
+        """Test get_labeled_for_export returns labeled documents."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [labeled_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_labeled_for_export()
+
+            assert len(result) == 1
+            assert result[0].status == "labeled"
+
+    def test_get_labeled_for_export_with_token(self, repo, labeled_document):
+        """Test get_labeled_for_export with admin_token filter."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [labeled_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_labeled_for_export(admin_token="token-123")
+
+            assert len(result) == 1
+
+    # ==========================================================================
+    # count_by_status() tests
+    # ==========================================================================
+
+    def test_count_by_status(self, repo):
+        """Test count_by_status returns status counts."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [
+                ("pending", 10),
+                ("labeled", 5),
+            ]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.count_by_status()
+
+            assert result == {"pending": 10, "labeled": 5}
+
+    # ==========================================================================
+    # get_by_ids() tests
+    # ==========================================================================
+
+    def test_get_by_ids(self, repo, sample_document):
+        """Test get_by_ids returns documents by IDs."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_by_ids([str(sample_document.document_id)])
+
+            assert len(result) == 1
+
+    # ==========================================================================
+    # get_for_training() tests
+    # ==========================================================================
+
+    def test_get_for_training_basic(self, repo, labeled_document):
+        """Test get_for_training with default parameters."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [labeled_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_for_training()
+
+            assert total == 1
+            assert len(results) == 1
+
+    def test_get_for_training_with_min_annotation_count(self, repo, labeled_document):
+        """Test get_for_training with min_annotation_count."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [labeled_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_for_training(min_annotation_count=3)
+
+            assert total == 1
+
+    def test_get_for_training_exclude_used(self, repo, labeled_document):
+        """Test get_for_training with exclude_used_in_training."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [labeled_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_for_training(exclude_used_in_training=True)
+
+            assert total == 1
+
+    def test_get_for_training_no_annotations(self, repo, labeled_document):
+        """Test get_for_training with has_annotations=False."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [labeled_document]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            results, total = repo.get_for_training(has_annotations=False)
+
+            assert total == 1
+
+    # ==========================================================================
+    # acquire_annotation_lock() tests
+    # ==========================================================================
+
+    def test_acquire_annotation_lock_success(self, repo, sample_document):
+        """Test acquire_annotation_lock when no lock exists."""
+        sample_document.annotation_lock_until = None
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.acquire_annotation_lock(str(sample_document.document_id))
+
+            assert result is not None
+            assert sample_document.annotation_lock_until is not None
+
+    def test_acquire_annotation_lock_fails_when_locked(self, repo, locked_document):
+        """Test acquire_annotation_lock fails when document is already locked."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = locked_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.acquire_annotation_lock(str(locked_document.document_id))
+
+            assert result is None
+
+    def test_acquire_annotation_lock_document_not_found(self, repo):
+        """Test acquire_annotation_lock when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.acquire_annotation_lock(str(uuid4()))
+
+            assert result is None
+
+    # ==========================================================================
+    # release_annotation_lock() tests
+    # ==========================================================================
+
+    def test_release_annotation_lock_success(self, repo, locked_document):
+        """Test release_annotation_lock releases the lock."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = locked_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.release_annotation_lock(str(locked_document.document_id))
+
+            assert result is not None
+            assert locked_document.annotation_lock_until is None
+
+    def test_release_annotation_lock_document_not_found(self, repo):
+        """Test release_annotation_lock when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.release_annotation_lock(str(uuid4()))
+
+            assert result is None
+
+    # ==========================================================================
+    # extend_annotation_lock() tests
+    # ==========================================================================
+
+    def test_extend_annotation_lock_success(self, repo, locked_document):
+        """Test extend_annotation_lock extends the lock."""
+        original_lock = locked_document.annotation_lock_until
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = locked_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.extend_annotation_lock(str(locked_document.document_id))
+
+            assert result is not None
+            assert locked_document.annotation_lock_until > original_lock
+
+    def test_extend_annotation_lock_fails_when_no_lock(self, repo, sample_document):
+        """Test extend_annotation_lock fails when no lock exists."""
+        sample_document.annotation_lock_until = None
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.extend_annotation_lock(str(sample_document.document_id))
+
+            assert result is None
+
+    def test_extend_annotation_lock_fails_when_expired(self, repo, expired_lock_document):
+        """Test extend_annotation_lock fails when lock is expired."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = expired_lock_document
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.extend_annotation_lock(str(expired_lock_document.document_id))
+
+            assert result is None
+
+    def test_extend_annotation_lock_document_not_found(self, repo):
+        """Test extend_annotation_lock when document not found."""
+        with patch("inference.data.repositories.document_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.extend_annotation_lock(str(uuid4()))
+
+            assert result is None

tests/data/repositories/test_model_version_repository.py

@@ -0,0 +1,582 @@
+"""
+Tests for ModelVersionRepository
+
+100% coverage tests for model version management.
+"""
+
+import pytest
+from datetime import datetime, timezone
+from unittest.mock import MagicMock, patch
+from uuid import uuid4, UUID
+
+from inference.data.admin_models import ModelVersion
+from inference.data.repositories.model_version_repository import ModelVersionRepository
+
+
+class TestModelVersionRepository:
+    """Tests for ModelVersionRepository."""
+
+    @pytest.fixture
+    def sample_model(self) -> ModelVersion:
+        """Create a sample model version for testing."""
+        return ModelVersion(
+            version_id=uuid4(),
+            version="v1.0.0",
+            name="Test Model",
+            description="A test model",
+            model_path="/path/to/model.pt",
+            status="ready",
+            is_active=False,
+            metrics_mAP=0.95,
+            metrics_precision=0.92,
+            metrics_recall=0.88,
+            document_count=100,
+            training_config={"epochs": 100},
+            file_size=1024000,
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def active_model(self) -> ModelVersion:
+        """Create an active model version for testing."""
+        return ModelVersion(
+            version_id=uuid4(),
+            version="v1.0.0",
+            name="Active Model",
+            model_path="/path/to/active_model.pt",
+            status="active",
+            is_active=True,
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def repo(self) -> ModelVersionRepository:
+        """Create a ModelVersionRepository instance."""
+        return ModelVersionRepository()
+
+    # =========================================================================
+    # create() tests
+    # =========================================================================
+
+    def test_create_returns_model(self, repo):
+        """Test create returns created model version."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                version="v1.0.0",
+                name="Test Model",
+                model_path="/path/to/model.pt",
+            )
+
+            mock_session.add.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_create_with_all_params(self, repo):
+        """Test create with all parameters."""
+        task_id = uuid4()
+        dataset_id = uuid4()
+        trained_at = datetime.now(timezone.utc)
+
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                version="v2.0.0",
+                name="Full Model",
+                model_path="/path/to/full_model.pt",
+                description="A complete model",
+                task_id=str(task_id),
+                dataset_id=str(dataset_id),
+                metrics_mAP=0.95,
+                metrics_precision=0.92,
+                metrics_recall=0.88,
+                document_count=500,
+                training_config={"epochs": 200},
+                file_size=2048000,
+                trained_at=trained_at,
+            )
+
+            added_model = mock_session.add.call_args[0][0]
+            assert added_model.version == "v2.0.0"
+            assert added_model.description == "A complete model"
+            assert added_model.task_id == task_id
+            assert added_model.dataset_id == dataset_id
+            assert added_model.metrics_mAP == 0.95
+
+    def test_create_with_uuid_objects(self, repo):
+        """Test create works with UUID objects."""
+        task_id = uuid4()
+        dataset_id = uuid4()
+
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(
+                version="v1.0.0",
+                name="Test Model",
+                model_path="/path/to/model.pt",
+                task_id=task_id,
+                dataset_id=dataset_id,
+            )
+
+            added_model = mock_session.add.call_args[0][0]
+            assert added_model.task_id == task_id
+            assert added_model.dataset_id == dataset_id
+
+    def test_create_without_optional_ids(self, repo):
+        """Test create without task_id and dataset_id."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(
+                version="v1.0.0",
+                name="Test Model",
+                model_path="/path/to/model.pt",
+            )
+
+            added_model = mock_session.add.call_args[0][0]
+            assert added_model.task_id is None
+            assert added_model.dataset_id is None
+
+    # =========================================================================
+    # get() tests
+    # =========================================================================
+
+    def test_get_returns_model(self, repo, sample_model):
+        """Test get returns model when exists."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(sample_model.version_id))
+
+            assert result is not None
+            assert result.name == "Test Model"
+            mock_session.expunge.assert_called_once()
+
+    def test_get_with_uuid(self, repo, sample_model):
+        """Test get works with UUID object."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(sample_model.version_id)
+
+            assert result is not None
+
+    def test_get_returns_none_when_not_found(self, repo):
+        """Test get returns None when model not found."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(uuid4()))
+
+            assert result is None
+            mock_session.expunge.assert_not_called()
+
+    # =========================================================================
+    # get_paginated() tests
+    # =========================================================================
+
+    def test_get_paginated_returns_models_and_total(self, repo, sample_model):
+        """Test get_paginated returns list of models and total count."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_model]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            models, total = repo.get_paginated()
+
+            assert len(models) == 1
+            assert total == 1
+
+    def test_get_paginated_with_status_filter(self, repo, sample_model):
+        """Test get_paginated filters by status."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_model]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            models, total = repo.get_paginated(status="ready")
+
+            assert len(models) == 1
+
+    def test_get_paginated_with_pagination(self, repo, sample_model):
+        """Test get_paginated with limit and offset."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 50
+            mock_session.exec.return_value.all.return_value = [sample_model]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            models, total = repo.get_paginated(limit=10, offset=20)
+
+            assert total == 50
+
+    def test_get_paginated_empty_results(self, repo):
+        """Test get_paginated with no results."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 0
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            models, total = repo.get_paginated()
+
+            assert models == []
+            assert total == 0
+
+    # =========================================================================
+    # get_active() tests
+    # =========================================================================
+
+    def test_get_active_returns_active_model(self, repo, active_model):
+        """Test get_active returns the active model."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.first.return_value = active_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_active()
+
+            assert result is not None
+            assert result.is_active is True
+            mock_session.expunge.assert_called_once()
+
+    def test_get_active_returns_none(self, repo):
+        """Test get_active returns None when no active model."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.first.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_active()
+
+            assert result is None
+            mock_session.expunge.assert_not_called()
+
+    # =========================================================================
+    # activate() tests
+    # =========================================================================
+
+    def test_activate_activates_model(self, repo, sample_model, active_model):
+        """Test activate sets model as active and deactivates others."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [active_model]
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.activate(str(sample_model.version_id))
+
+            assert result is not None
+            assert sample_model.is_active is True
+            assert sample_model.status == "active"
+            assert active_model.is_active is False
+            assert active_model.status == "inactive"
+
+    def test_activate_with_uuid(self, repo, sample_model):
+        """Test activate works with UUID object."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.activate(sample_model.version_id)
+
+            assert result is not None
+            assert sample_model.is_active is True
+
+    def test_activate_returns_none_when_not_found(self, repo):
+        """Test activate returns None when model not found."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.activate(str(uuid4()))
+
+            assert result is None
+
+    def test_activate_sets_activated_at(self, repo, sample_model):
+        """Test activate sets activated_at timestamp."""
+        sample_model.activated_at = None
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.activate(str(sample_model.version_id))
+
+            assert sample_model.activated_at is not None
+
+    # =========================================================================
+    # deactivate() tests
+    # =========================================================================
+
+    def test_deactivate_deactivates_model(self, repo, active_model):
+        """Test deactivate sets model as inactive."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = active_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.deactivate(str(active_model.version_id))
+
+            assert result is not None
+            assert active_model.is_active is False
+            assert active_model.status == "inactive"
+            mock_session.commit.assert_called_once()
+
+    def test_deactivate_with_uuid(self, repo, active_model):
+        """Test deactivate works with UUID object."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = active_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.deactivate(active_model.version_id)
+
+            assert result is not None
+
+    def test_deactivate_returns_none_when_not_found(self, repo):
+        """Test deactivate returns None when model not found."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.deactivate(str(uuid4()))
+
+            assert result is None
+
+    # =========================================================================
+    # update() tests
+    # =========================================================================
+
+    def test_update_updates_model(self, repo, sample_model):
+        """Test update updates model metadata."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(
+                str(sample_model.version_id),
+                name="Updated Model",
+            )
+
+            assert result is not None
+            assert sample_model.name == "Updated Model"
+            mock_session.commit.assert_called_once()
+
+    def test_update_all_fields(self, repo, sample_model):
+        """Test update can update all fields."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(
+                str(sample_model.version_id),
+                name="New Name",
+                description="New Description",
+                status="archived",
+            )
+
+            assert sample_model.name == "New Name"
+            assert sample_model.description == "New Description"
+            assert sample_model.status == "archived"
+
+    def test_update_with_uuid(self, repo, sample_model):
+        """Test update works with UUID object."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(sample_model.version_id, name="Updated")
+
+            assert result is not None
+
+    def test_update_returns_none_when_not_found(self, repo):
+        """Test update returns None when model not found."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(str(uuid4()), name="New Name")
+
+            assert result is None
+
+    def test_update_partial_fields(self, repo, sample_model):
+        """Test update only updates provided fields."""
+        original_name = sample_model.name
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.update(
+                str(sample_model.version_id),
+                description="Only description changed",
+            )
+
+            assert sample_model.name == original_name
+            assert sample_model.description == "Only description changed"
+
+    # =========================================================================
+    # archive() tests
+    # =========================================================================
+
+    def test_archive_archives_model(self, repo, sample_model):
+        """Test archive sets model status to archived."""
+        sample_model.is_active = False
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.archive(str(sample_model.version_id))
+
+            assert result is not None
+            assert sample_model.status == "archived"
+            mock_session.commit.assert_called_once()
+
+    def test_archive_with_uuid(self, repo, sample_model):
+        """Test archive works with UUID object."""
+        sample_model.is_active = False
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.archive(sample_model.version_id)
+
+            assert result is not None
+
+    def test_archive_returns_none_when_not_found(self, repo):
+        """Test archive returns None when model not found."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.archive(str(uuid4()))
+
+            assert result is None
+
+    def test_archive_returns_none_when_active(self, repo, active_model):
+        """Test archive returns None when model is active."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = active_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.archive(str(active_model.version_id))
+
+            assert result is None
+
+    # =========================================================================
+    # delete() tests
+    # =========================================================================
+
+    def test_delete_returns_true(self, repo, sample_model):
+        """Test delete returns True when model exists and not active."""
+        sample_model.is_active = False
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(sample_model.version_id))
+
+            assert result is True
+            mock_session.delete.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_delete_with_uuid(self, repo, sample_model):
+        """Test delete works with UUID object."""
+        sample_model.is_active = False
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(sample_model.version_id)
+
+            assert result is True
+
+    def test_delete_returns_false_when_not_found(self, repo):
+        """Test delete returns False when model not found."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(uuid4()))
+
+            assert result is False
+            mock_session.delete.assert_not_called()
+
+    def test_delete_returns_false_when_active(self, repo, active_model):
+        """Test delete returns False when model is active."""
+        with patch("inference.data.repositories.model_version_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = active_model
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.delete(str(active_model.version_id))
+
+            assert result is False
+            mock_session.delete.assert_not_called()

tests/data/repositories/test_token_repository.py

@@ -0,0 +1,199 @@
+"""
+Tests for TokenRepository
+
+TDD tests for admin token management.
+"""
+
+import pytest
+from datetime import datetime, timedelta, timezone
+from unittest.mock import MagicMock, patch
+
+from inference.data.admin_models import AdminToken
+from inference.data.repositories.token_repository import TokenRepository
+
+
+class TestTokenRepository:
+    """Tests for TokenRepository."""
+
+    @pytest.fixture
+    def sample_token(self) -> AdminToken:
+        """Create a sample token for testing."""
+        return AdminToken(
+            token="test-token-123",
+            name="Test Token",
+            is_active=True,
+            created_at=datetime.now(timezone.utc),
+            last_used_at=None,
+            expires_at=None,
+        )
+
+    @pytest.fixture
+    def expired_token(self) -> AdminToken:
+        """Create an expired token."""
+        return AdminToken(
+            token="expired-token",
+            name="Expired Token",
+            is_active=True,
+            created_at=datetime.now(timezone.utc) - timedelta(days=30),
+            expires_at=datetime.now(timezone.utc) - timedelta(days=1),
+        )
+
+    @pytest.fixture
+    def inactive_token(self) -> AdminToken:
+        """Create an inactive token."""
+        return AdminToken(
+            token="inactive-token",
+            name="Inactive Token",
+            is_active=False,
+            created_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def repo(self) -> TokenRepository:
+        """Create a TokenRepository instance."""
+        return TokenRepository()
+
+    def test_is_valid_returns_true_for_active_token(self, repo, sample_token):
+        """Test is_valid returns True for an active, non-expired token."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_token
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.is_valid("test-token-123")
+
+            assert result is True
+            mock_session.get.assert_called_once_with(AdminToken, "test-token-123")
+
+    def test_is_valid_returns_false_for_nonexistent_token(self, repo):
+        """Test is_valid returns False for a non-existent token."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.is_valid("nonexistent-token")
+
+            assert result is False
+
+    def test_is_valid_returns_false_for_inactive_token(self, repo, inactive_token):
+        """Test is_valid returns False for an inactive token."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = inactive_token
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.is_valid("inactive-token")
+
+            assert result is False
+
+    def test_is_valid_returns_false_for_expired_token(self, repo, expired_token):
+        """Test is_valid returns False for an expired token."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = expired_token
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.is_valid("expired-token")
+
+            assert result is False
+
+    def test_get_returns_token_when_exists(self, repo, sample_token):
+        """Test get returns token when it exists."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_token
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get("test-token-123")
+
+            assert result is not None
+            assert result.token == "test-token-123"
+            assert result.name == "Test Token"
+            mock_session.expunge.assert_called_once_with(sample_token)
+
+    def test_get_returns_none_when_not_exists(self, repo):
+        """Test get returns None when token doesn't exist."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get("nonexistent-token")
+
+            assert result is None
+
+    def test_create_new_token(self, repo):
+        """Test creating a new token."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None  # Token doesn't exist
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create("new-token", "New Token", expires_at=None)
+
+            mock_session.add.assert_called_once()
+            added_token = mock_session.add.call_args[0][0]
+            assert isinstance(added_token, AdminToken)
+            assert added_token.token == "new-token"
+            assert added_token.name == "New Token"
+
+    def test_create_updates_existing_token(self, repo, sample_token):
+        """Test create updates an existing token."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_token
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create("test-token-123", "Updated Name", expires_at=None)
+
+            mock_session.add.assert_called_once_with(sample_token)
+            assert sample_token.name == "Updated Name"
+            assert sample_token.is_active is True
+
+    def test_update_usage(self, repo, sample_token):
+        """Test updating token last_used_at timestamp."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_token
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_usage("test-token-123")
+
+            assert sample_token.last_used_at is not None
+            mock_session.add.assert_called_once_with(sample_token)
+
+    def test_deactivate_returns_true_when_token_exists(self, repo, sample_token):
+        """Test deactivate returns True when token exists."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_token
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.deactivate("test-token-123")
+
+            assert result is True
+            assert sample_token.is_active is False
+            mock_session.add.assert_called_once_with(sample_token)
+
+    def test_deactivate_returns_false_when_token_not_exists(self, repo):
+        """Test deactivate returns False when token doesn't exist."""
+        with patch("inference.data.repositories.base.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.deactivate("nonexistent-token")
+
+            assert result is False

tests/data/repositories/test_training_task_repository.py

@@ -0,0 +1,615 @@
+"""
+Tests for TrainingTaskRepository
+
+100% coverage tests for training task management.
+"""
+
+import pytest
+from datetime import datetime, timezone, timedelta
+from unittest.mock import MagicMock, patch
+from uuid import uuid4, UUID
+
+from inference.data.admin_models import TrainingTask, TrainingLog, TrainingDocumentLink
+from inference.data.repositories.training_task_repository import TrainingTaskRepository
+
+
+class TestTrainingTaskRepository:
+    """Tests for TrainingTaskRepository."""
+
+    @pytest.fixture
+    def sample_task(self) -> TrainingTask:
+        """Create a sample training task for testing."""
+        return TrainingTask(
+            task_id=uuid4(),
+            admin_token="admin-token",
+            name="Test Training Task",
+            task_type="train",
+            description="A test training task",
+            status="pending",
+            config={"epochs": 100, "batch_size": 16},
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def sample_log(self) -> TrainingLog:
+        """Create a sample training log for testing."""
+        return TrainingLog(
+            log_id=uuid4(),
+            task_id=uuid4(),
+            level="INFO",
+            message="Training started",
+            details={"epoch": 1},
+            created_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def sample_link(self) -> TrainingDocumentLink:
+        """Create a sample training document link for testing."""
+        return TrainingDocumentLink(
+            link_id=uuid4(),
+            task_id=uuid4(),
+            document_id=uuid4(),
+            annotation_snapshot={"annotations": []},
+            created_at=datetime.now(timezone.utc),
+        )
+
+    @pytest.fixture
+    def repo(self) -> TrainingTaskRepository:
+        """Create a TrainingTaskRepository instance."""
+        return TrainingTaskRepository()
+
+    # =========================================================================
+    # create() tests
+    # =========================================================================
+
+    def test_create_returns_task_id(self, repo):
+        """Test create returns task ID."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                admin_token="admin-token",
+                name="Test Task",
+            )
+
+            assert result is not None
+            mock_session.add.assert_called_once()
+            mock_session.flush.assert_called_once()
+
+    def test_create_with_all_params(self, repo):
+        """Test create with all parameters."""
+        scheduled_time = datetime.now(timezone.utc) + timedelta(hours=1)
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.create(
+                admin_token="admin-token",
+                name="Test Task",
+                task_type="finetune",
+                description="Full test",
+                config={"epochs": 50},
+                scheduled_at=scheduled_time,
+                cron_expression="0 0 * * *",
+                is_recurring=True,
+                dataset_id=str(uuid4()),
+            )
+
+            assert result is not None
+            added_task = mock_session.add.call_args[0][0]
+            assert added_task.task_type == "finetune"
+            assert added_task.description == "Full test"
+            assert added_task.is_recurring is True
+            assert added_task.status == "scheduled"  # because scheduled_at is set
+
+    def test_create_pending_status_when_not_scheduled(self, repo):
+        """Test create sets pending status when no scheduled_at."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(
+                admin_token="admin-token",
+                name="Test Task",
+            )
+
+            added_task = mock_session.add.call_args[0][0]
+            assert added_task.status == "pending"
+
+    def test_create_scheduled_status_when_scheduled(self, repo):
+        """Test create sets scheduled status when scheduled_at is provided."""
+        scheduled_time = datetime.now(timezone.utc) + timedelta(hours=1)
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.create(
+                admin_token="admin-token",
+                name="Test Task",
+                scheduled_at=scheduled_time,
+            )
+
+            added_task = mock_session.add.call_args[0][0]
+            assert added_task.status == "scheduled"
+
+    # =========================================================================
+    # get() tests
+    # =========================================================================
+
+    def test_get_returns_task(self, repo, sample_task):
+        """Test get returns task when exists."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(sample_task.task_id))
+
+            assert result is not None
+            assert result.name == "Test Training Task"
+            mock_session.expunge.assert_called_once()
+
+    def test_get_returns_none_when_not_found(self, repo):
+        """Test get returns None when task not found."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get(str(uuid4()))
+
+            assert result is None
+            mock_session.expunge.assert_not_called()
+
+    # =========================================================================
+    # get_by_token() tests
+    # =========================================================================
+
+    def test_get_by_token_returns_task(self, repo, sample_task):
+        """Test get_by_token returns task (delegates to get)."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_by_token(str(sample_task.task_id), "admin-token")
+
+            assert result is not None
+
+    def test_get_by_token_without_token_param(self, repo, sample_task):
+        """Test get_by_token works without token parameter."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_by_token(str(sample_task.task_id))
+
+            assert result is not None
+
+    # =========================================================================
+    # get_paginated() tests
+    # =========================================================================
+
+    def test_get_paginated_returns_tasks_and_total(self, repo, sample_task):
+        """Test get_paginated returns list of tasks and total count."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_task]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            tasks, total = repo.get_paginated()
+
+            assert len(tasks) == 1
+            assert total == 1
+
+    def test_get_paginated_with_status_filter(self, repo, sample_task):
+        """Test get_paginated filters by status."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 1
+            mock_session.exec.return_value.all.return_value = [sample_task]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            tasks, total = repo.get_paginated(status="pending")
+
+            assert len(tasks) == 1
+
+    def test_get_paginated_with_pagination(self, repo, sample_task):
+        """Test get_paginated with limit and offset."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 50
+            mock_session.exec.return_value.all.return_value = [sample_task]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            tasks, total = repo.get_paginated(limit=10, offset=20)
+
+            assert total == 50
+
+    def test_get_paginated_empty_results(self, repo):
+        """Test get_paginated with no results."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.one.return_value = 0
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            tasks, total = repo.get_paginated()
+
+            assert tasks == []
+            assert total == 0
+
+    # =========================================================================
+    # get_pending() tests
+    # =========================================================================
+
+    def test_get_pending_returns_pending_tasks(self, repo, sample_task):
+        """Test get_pending returns pending and scheduled tasks."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_task]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_pending()
+
+            assert len(result) == 1
+
+    def test_get_pending_returns_empty_list(self, repo):
+        """Test get_pending returns empty list when no pending tasks."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_pending()
+
+            assert result == []
+
+    # =========================================================================
+    # update_status() tests
+    # =========================================================================
+
+    def test_update_status_updates_task(self, repo, sample_task):
+        """Test update_status updates task status."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(sample_task.task_id), "running")
+
+            assert sample_task.status == "running"
+
+    def test_update_status_sets_started_at_for_running(self, repo, sample_task):
+        """Test update_status sets started_at when status is running."""
+        sample_task.started_at = None
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(sample_task.task_id), "running")
+
+            assert sample_task.started_at is not None
+
+    def test_update_status_sets_completed_at_for_completed(self, repo, sample_task):
+        """Test update_status sets completed_at when status is completed."""
+        sample_task.completed_at = None
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(sample_task.task_id), "completed")
+
+            assert sample_task.completed_at is not None
+
+    def test_update_status_sets_completed_at_for_failed(self, repo, sample_task):
+        """Test update_status sets completed_at when status is failed."""
+        sample_task.completed_at = None
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(sample_task.task_id), "failed", error_message="Error occurred")
+
+            assert sample_task.completed_at is not None
+            assert sample_task.error_message == "Error occurred"
+
+    def test_update_status_with_result_metrics(self, repo, sample_task):
+        """Test update_status with result metrics."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(
+                str(sample_task.task_id),
+                "completed",
+                result_metrics={"mAP": 0.95},
+            )
+
+            assert sample_task.result_metrics == {"mAP": 0.95}
+
+    def test_update_status_with_model_path(self, repo, sample_task):
+        """Test update_status with model path."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(
+                str(sample_task.task_id),
+                "completed",
+                model_path="/path/to/model.pt",
+            )
+
+            assert sample_task.model_path == "/path/to/model.pt"
+
+    def test_update_status_not_found(self, repo):
+        """Test update_status does nothing when task not found."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.update_status(str(uuid4()), "running")
+
+            mock_session.add.assert_not_called()
+
+    # =========================================================================
+    # cancel() tests
+    # =========================================================================
+
+    def test_cancel_returns_true_for_pending(self, repo, sample_task):
+        """Test cancel returns True for pending task."""
+        sample_task.status = "pending"
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.cancel(str(sample_task.task_id))
+
+            assert result is True
+            assert sample_task.status == "cancelled"
+
+    def test_cancel_returns_true_for_scheduled(self, repo, sample_task):
+        """Test cancel returns True for scheduled task."""
+        sample_task.status = "scheduled"
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.cancel(str(sample_task.task_id))
+
+            assert result is True
+            assert sample_task.status == "cancelled"
+
+    def test_cancel_returns_false_for_running(self, repo, sample_task):
+        """Test cancel returns False for running task."""
+        sample_task.status = "running"
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = sample_task
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.cancel(str(sample_task.task_id))
+
+            assert result is False
+
+    def test_cancel_returns_false_when_not_found(self, repo):
+        """Test cancel returns False when task not found."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.get.return_value = None
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.cancel(str(uuid4()))
+
+            assert result is False
+
+    # =========================================================================
+    # add_log() tests
+    # =========================================================================
+
+    def test_add_log_creates_log_entry(self, repo):
+        """Test add_log creates a log entry."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.add_log(
+                task_id=str(uuid4()),
+                level="INFO",
+                message="Training started",
+            )
+
+            mock_session.add.assert_called_once()
+            added_log = mock_session.add.call_args[0][0]
+            assert added_log.level == "INFO"
+            assert added_log.message == "Training started"
+
+    def test_add_log_with_details(self, repo):
+        """Test add_log with details."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            repo.add_log(
+                task_id=str(uuid4()),
+                level="DEBUG",
+                message="Epoch complete",
+                details={"epoch": 5, "loss": 0.05},
+            )
+
+            added_log = mock_session.add.call_args[0][0]
+            assert added_log.details == {"epoch": 5, "loss": 0.05}
+
+    # =========================================================================
+    # get_logs() tests
+    # =========================================================================
+
+    def test_get_logs_returns_list(self, repo, sample_log):
+        """Test get_logs returns list of logs."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_log]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_logs(str(sample_log.task_id))
+
+            assert len(result) == 1
+            assert result[0].level == "INFO"
+
+    def test_get_logs_with_pagination(self, repo, sample_log):
+        """Test get_logs with limit and offset."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_log]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_logs(str(sample_log.task_id), limit=50, offset=10)
+
+            assert len(result) == 1
+
+    def test_get_logs_returns_empty_list(self, repo):
+        """Test get_logs returns empty list when no logs."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_logs(str(uuid4()))
+
+            assert result == []
+
+    # =========================================================================
+    # create_document_link() tests
+    # =========================================================================
+
+    def test_create_document_link_returns_link(self, repo):
+        """Test create_document_link returns created link."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            task_id = uuid4()
+            document_id = uuid4()
+            result = repo.create_document_link(
+                task_id=task_id,
+                document_id=document_id,
+            )
+
+            mock_session.add.assert_called_once()
+            mock_session.commit.assert_called_once()
+
+    def test_create_document_link_with_snapshot(self, repo):
+        """Test create_document_link with annotation snapshot."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            snapshot = {"annotations": [{"class_name": "invoice_number"}]}
+            repo.create_document_link(
+                task_id=uuid4(),
+                document_id=uuid4(),
+                annotation_snapshot=snapshot,
+            )
+
+            added_link = mock_session.add.call_args[0][0]
+            assert added_link.annotation_snapshot == snapshot
+
+    # =========================================================================
+    # get_document_links() tests
+    # =========================================================================
+
+    def test_get_document_links_returns_list(self, repo, sample_link):
+        """Test get_document_links returns list of links."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_link]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_document_links(sample_link.task_id)
+
+            assert len(result) == 1
+
+    def test_get_document_links_returns_empty_list(self, repo):
+        """Test get_document_links returns empty list when no links."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_document_links(uuid4())
+
+            assert result == []
+
+    # =========================================================================
+    # get_document_training_tasks() tests
+    # =========================================================================
+
+    def test_get_document_training_tasks_returns_list(self, repo, sample_link):
+        """Test get_document_training_tasks returns list of links."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = [sample_link]
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_document_training_tasks(sample_link.document_id)
+
+            assert len(result) == 1
+
+    def test_get_document_training_tasks_returns_empty_list(self, repo):
+        """Test get_document_training_tasks returns empty list when no links."""
+        with patch("inference.data.repositories.training_task_repository.get_session_context") as mock_ctx:
+            mock_session = MagicMock()
+            mock_session.exec.return_value.all.return_value = []
+            mock_ctx.return_value.__enter__ = MagicMock(return_value=mock_session)
+            mock_ctx.return_value.__exit__ = MagicMock(return_value=False)
+
+            result = repo.get_document_training_tasks(uuid4())
+
+            assert result == []

tests/inference/test_field_extractor.py

@@ -12,6 +12,15 @@ Tests field normalization functions:
 
 import pytest
 from inference.pipeline.field_extractor import FieldExtractor
+from inference.pipeline.normalizers import (
+    InvoiceNumberNormalizer,
+    OcrNumberNormalizer,
+    BankgiroNormalizer,
+    PlusgiroNormalizer,
+    AmountNormalizer,
+    DateNormalizer,
+    SupplierOrgNumberNormalizer,
+)
 
 
 class TestFieldExtractorInit:
@@ -43,81 +52,81 @@ class TestNormalizeInvoiceNumber:
     """Tests for invoice number normalization."""
 
     @pytest.fixture
-    def extractor(self):
-        return FieldExtractor()
+    def normalizer(self):
+        return InvoiceNumberNormalizer()
 
-    def test_alphanumeric_invoice_number(self, extractor):
+    def test_alphanumeric_invoice_number(self, normalizer):
         """Test alphanumeric invoice number like A3861."""
-        result, is_valid, error = extractor._normalize_invoice_number("Fakturanummer: A3861")
-        assert result == 'A3861'
-        assert is_valid is True
+        result = normalizer.normalize("Fakturanummer: A3861")
+        assert result.value == 'A3861'
+        assert result.is_valid is True
 
-    def test_prefix_invoice_number(self, extractor):
+    def test_prefix_invoice_number(self, normalizer):
         """Test invoice number with prefix like INV12345."""
-        result, is_valid, error = extractor._normalize_invoice_number("Invoice INV12345")
-        assert result is not None
-        assert 'INV' in result or '12345' in result
+        result = normalizer.normalize("Invoice INV12345")
+        assert result.value is not None
+        assert 'INV' in result.value or '12345' in result.value
 
-    def test_numeric_invoice_number(self, extractor):
+    def test_numeric_invoice_number(self, normalizer):
         """Test pure numeric invoice number."""
-        result, is_valid, error = extractor._normalize_invoice_number("Invoice: 12345678")
-        assert result is not None
-        assert result.isdigit()
+        result = normalizer.normalize("Invoice: 12345678")
+        assert result.value is not None
+        assert result.value.isdigit()
 
-    def test_year_prefixed_invoice_number(self, extractor):
+    def test_year_prefixed_invoice_number(self, normalizer):
         """Test invoice number with year prefix like 2024-001."""
-        result, is_valid, error = extractor._normalize_invoice_number("Faktura 2024-12345")
-        assert result is not None
-        assert '2024' in result
+        result = normalizer.normalize("Faktura 2024-12345")
+        assert result.value is not None
+        assert '2024' in result.value
 
-    def test_avoid_long_ocr_sequence(self, extractor):
+    def test_avoid_long_ocr_sequence(self, normalizer):
         """Test that long OCR-like sequences are avoided."""
         # When text contains both short invoice number and long OCR sequence
         text = "Fakturanummer: A3861 OCR: 310196187399952763290708"
-        result, is_valid, error = extractor._normalize_invoice_number(text)
+        result = normalizer.normalize(text)
         # Should prefer the shorter alphanumeric pattern
-        assert result == 'A3861'
+        assert result.value == 'A3861'
 
-    def test_empty_string(self, extractor):
+    def test_empty_string(self, normalizer):
         """Test empty string input."""
-        result, is_valid, error = extractor._normalize_invoice_number("")
-        assert result is None or is_valid is False
+        result = normalizer.normalize("")
+        assert result.value is None or result.is_valid is False
 
 
 class TestNormalizeBankgiro:
     """Tests for Bankgiro normalization."""
 
     @pytest.fixture
-    def extractor(self):
-        return FieldExtractor()
+    def normalizer(self):
+        return BankgiroNormalizer()
 
-    def test_standard_7_digit_format(self, extractor):
+    def test_standard_7_digit_format(self, normalizer):
         """Test 7-digit Bankgiro XXX-XXXX."""
-        result, is_valid, error = extractor._normalize_bankgiro("Bankgiro: 782-1713")
-        assert result == '782-1713'
-        assert is_valid is True
+        result = normalizer.normalize("Bankgiro: 782-1713")
+        assert result.value == '782-1713'
+        assert result.is_valid is True
 
-    def test_standard_8_digit_format(self, extractor):
+    def test_standard_8_digit_format(self, normalizer):
         """Test 8-digit Bankgiro XXXX-XXXX."""
-        result, is_valid, error = extractor._normalize_bankgiro("BG 5393-9484")
-        assert result == '5393-9484'
-        assert is_valid is True
+        result = normalizer.normalize("BG 5393-9484")
+        assert result.value == '5393-9484'
+        assert result.is_valid is True
 
-    def test_without_dash(self, extractor):
+    def test_without_dash(self, normalizer):
         """Test Bankgiro without dash."""
-        result, is_valid, error = extractor._normalize_bankgiro("Bankgiro 7821713")
-        assert result is not None
+        result = normalizer.normalize("Bankgiro 7821713")
+        assert result.value is not None
         # Should be formatted with dash
 
-    def test_with_spaces(self, extractor):
+    def test_with_spaces(self, normalizer):
         """Test Bankgiro with spaces - may not parse if spaces break the pattern."""
-        result, is_valid, error = extractor._normalize_bankgiro("BG: 782 1713")
+        result = normalizer.normalize("BG: 782 1713")
         # Spaces in the middle might cause parsing issues - that's acceptable
         # The test passes if it doesn't crash
 
-    def test_invalid_bankgiro(self, extractor):
+    def test_invalid_bankgiro(self, normalizer):
         """Test invalid Bankgiro (too short)."""
-        result, is_valid, error = extractor._normalize_bankgiro("BG: 123")
+        result = normalizer.normalize("BG: 123")
         # Should fail or return None
 
 
@@ -125,28 +134,32 @@ class TestNormalizePlusgiro:
     """Tests for Plusgiro normalization."""
 
     @pytest.fixture
-    def extractor(self):
-        return FieldExtractor()
+    def normalizer(self):
+        return PlusgiroNormalizer()
 
-    def test_standard_format(self, extractor):
+    @pytest.fixture
+    def bg_normalizer(self):
+        return BankgiroNormalizer()
+
+    def test_standard_format(self, normalizer):
         """Test standard Plusgiro format XXXXXXX-X."""
-        result, is_valid, error = extractor._normalize_plusgiro("Plusgiro: 1234567-8")
-        assert result is not None
-        assert '-' in result
+        result = normalizer.normalize("Plusgiro: 1234567-8")
+        assert result.value is not None
+        assert '-' in result.value
 
-    def test_without_dash(self, extractor):
+    def test_without_dash(self, normalizer):
         """Test Plusgiro without dash."""
-        result, is_valid, error = extractor._normalize_plusgiro("PG 12345678")
-        assert result is not None
+        result = normalizer.normalize("PG 12345678")
+        assert result.value is not None
 
-    def test_distinguish_from_bankgiro(self, extractor):
+    def test_distinguish_from_bankgiro(self, normalizer, bg_normalizer):
         """Test that Plusgiro is distinguished from Bankgiro by format."""
         # Plusgiro has 1 digit after dash, Bankgiro has 4
         pg_text = "4809603-6"  # Plusgiro format
         bg_text = "782-1713"  # Bankgiro format
 
-        pg_result, _, _ = extractor._normalize_plusgiro(pg_text)
-        bg_result, _, _ = extractor._normalize_bankgiro(bg_text)
+        pg_result = normalizer.normalize(pg_text)
+        bg_result = bg_normalizer.normalize(bg_text)
 
         # Both should succeed in their respective normalizations
 
@@ -155,89 +168,89 @@ class TestNormalizeAmount:
     """Tests for Amount normalization."""
 
     @pytest.fixture
-    def extractor(self):
-        return FieldExtractor()
+    def normalizer(self):
+        return AmountNormalizer()
 
-    def test_swedish_format_comma(self, extractor):
+    def test_swedish_format_comma(self, normalizer):
         """Test Swedish format with comma: 11 699,00."""
-        result, is_valid, error = extractor._normalize_amount("11 699,00 SEK")
-        assert result is not None
-        assert is_valid is True
+        result = normalizer.normalize("11 699,00 SEK")
+        assert result.value is not None
+        assert result.is_valid is True
 
-    def test_integer_amount(self, extractor):
+    def test_integer_amount(self, normalizer):
         """Test integer amount without decimals."""
-        result, is_valid, error = extractor._normalize_amount("Amount: 11699")
-        assert result is not None
+        result = normalizer.normalize("Amount: 11699")
+        assert result.value is not None
 
-    def test_with_currency(self, extractor):
+    def test_with_currency(self, normalizer):
         """Test amount with currency symbol."""
-        result, is_valid, error = extractor._normalize_amount("SEK 11 699,00")
-        assert result is not None
+        result = normalizer.normalize("SEK 11 699,00")
+        assert result.value is not None
 
-    def test_large_amount(self, extractor):
+    def test_large_amount(self, normalizer):
         """Test large amount with thousand separators."""
-        result, is_valid, error = extractor._normalize_amount("1 234 567,89")
-        assert result is not None
+        result = normalizer.normalize("1 234 567,89")
+        assert result.value is not None
 
 
 class TestNormalizeOCR:
     """Tests for OCR number normalization."""
 
     @pytest.fixture
-    def extractor(self):
-        return FieldExtractor()
+    def normalizer(self):
+        return OcrNumberNormalizer()
 
-    def test_standard_ocr(self, extractor):
+    def test_standard_ocr(self, normalizer):
         """Test standard OCR number."""
-        result, is_valid, error = extractor._normalize_ocr_number("OCR: 310196187399952")
-        assert result == '310196187399952'
-        assert is_valid is True
+        result = normalizer.normalize("OCR: 310196187399952")
+        assert result.value == '310196187399952'
+        assert result.is_valid is True
 
-    def test_ocr_with_spaces(self, extractor):
+    def test_ocr_with_spaces(self, normalizer):
         """Test OCR number with spaces."""
-        result, is_valid, error = extractor._normalize_ocr_number("3101 9618 7399 952")
-        assert result is not None
-        assert ' ' not in result  # Spaces should be removed
+        result = normalizer.normalize("3101 9618 7399 952")
+        assert result.value is not None
+        assert ' ' not in result.value  # Spaces should be removed
 
-    def test_short_ocr_invalid(self, extractor):
+    def test_short_ocr_invalid(self, normalizer):
         """Test that too short OCR is invalid."""
-        result, is_valid, error = extractor._normalize_ocr_number("123")
-        assert is_valid is False
+        result = normalizer.normalize("123")
+        assert result.is_valid is False
 
 
 class TestNormalizeDate:
     """Tests for date normalization."""
 
     @pytest.fixture
-    def extractor(self):
-        return FieldExtractor()
+    def normalizer(self):
+        return DateNormalizer()
 
-    def test_iso_format(self, extractor):
+    def test_iso_format(self, normalizer):
         """Test ISO date format YYYY-MM-DD."""
-        result, is_valid, error = extractor._normalize_date("2026-01-31")
-        assert result == '2026-01-31'
-        assert is_valid is True
+        result = normalizer.normalize("2026-01-31")
+        assert result.value == '2026-01-31'
+        assert result.is_valid is True
 
-    def test_swedish_format(self, extractor):
+    def test_swedish_format(self, normalizer):
         """Test Swedish format with dots: 31.01.2026."""
-        result, is_valid, error = extractor._normalize_date("31.01.2026")
-        assert result is not None
-        assert is_valid is True
+        result = normalizer.normalize("31.01.2026")
+        assert result.value is not None
+        assert result.is_valid is True
 
-    def test_slash_format(self, extractor):
+    def test_slash_format(self, normalizer):
         """Test slash format: 31/01/2026."""
-        result, is_valid, error = extractor._normalize_date("31/01/2026")
-        assert result is not None
+        result = normalizer.normalize("31/01/2026")
+        assert result.value is not None
 
-    def test_compact_format(self, extractor):
+    def test_compact_format(self, normalizer):
         """Test compact format: 20260131."""
-        result, is_valid, error = extractor._normalize_date("20260131")
-        assert result is not None
+        result = normalizer.normalize("20260131")
+        assert result.value is not None
 
-    def test_invalid_date(self, extractor):
+    def test_invalid_date(self, normalizer):
         """Test invalid date."""
-        result, is_valid, error = extractor._normalize_date("not a date")
-        assert is_valid is False
+        result = normalizer.normalize("not a date")
+        assert result.is_valid is False
 
 
 class TestNormalizePaymentLine:
@@ -348,20 +361,20 @@ class TestNormalizeSupplierOrgNumber:
     """Tests for supplier organization number normalization."""
 
     @pytest.fixture
-    def extractor(self):
-        return FieldExtractor()
+    def normalizer(self):
+        return SupplierOrgNumberNormalizer()
 
-    def test_standard_format(self, extractor):
+    def test_standard_format(self, normalizer):
         """Test standard format NNNNNN-NNNN."""
-        result, is_valid, error = extractor._normalize_supplier_org_number("Org.nr 516406-1102")
-        assert result == '516406-1102'
-        assert is_valid is True
+        result = normalizer.normalize("Org.nr 516406-1102")
+        assert result.value == '516406-1102'
+        assert result.is_valid is True
 
-    def test_vat_number_format(self, extractor):
+    def test_vat_number_format(self, normalizer):
         """Test VAT number format SE + 10 digits + 01."""
-        result, is_valid, error = extractor._normalize_supplier_org_number("Momsreg.nr SE556123456701")
-        assert result is not None
-        assert '-' in result
+        result = normalizer.normalize("Momsreg.nr SE556123456701")
+        assert result.value is not None
+        assert '-' in result.value
 
 
 class TestNormalizeAndValidateDispatch:

tests/inference/test_normalizers.py

@@ -0,0 +1,768 @@
+"""
+Tests for Inference Pipeline Normalizers
+
+These normalizers extract and validate field values from OCR text.
+They are different from shared/normalize/normalizers which generate
+matching variants from known values.
+"""
+
+from unittest.mock import patch
+import pytest
+from inference.pipeline.normalizers import (
+    NormalizationResult,
+    InvoiceNumberNormalizer,
+    OcrNumberNormalizer,
+    BankgiroNormalizer,
+    PlusgiroNormalizer,
+    AmountNormalizer,
+    EnhancedAmountNormalizer,
+    DateNormalizer,
+    EnhancedDateNormalizer,
+    SupplierOrgNumberNormalizer,
+    create_normalizer_registry,
+)
+
+
+class TestNormalizationResult:
+    """Tests for NormalizationResult dataclass."""
+
+    def test_success(self):
+        result = NormalizationResult.success("123")
+        assert result.value == "123"
+        assert result.is_valid is True
+        assert result.error is None
+
+    def test_success_with_warning(self):
+        result = NormalizationResult.success_with_warning("123", "Warning message")
+        assert result.value == "123"
+        assert result.is_valid is True
+        assert result.error == "Warning message"
+
+    def test_failure(self):
+        result = NormalizationResult.failure("Error message")
+        assert result.value is None
+        assert result.is_valid is False
+        assert result.error == "Error message"
+
+    def test_to_tuple(self):
+        result = NormalizationResult.success("123")
+        value, is_valid, error = result.to_tuple()
+        assert value == "123"
+        assert is_valid is True
+        assert error is None
+
+
+class TestInvoiceNumberNormalizer:
+    """Tests for InvoiceNumberNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return InvoiceNumberNormalizer()
+
+    def test_field_name(self, normalizer):
+        assert normalizer.field_name == "InvoiceNumber"
+
+    def test_alphanumeric(self, normalizer):
+        result = normalizer.normalize("A3861")
+        assert result.value == "A3861"
+        assert result.is_valid is True
+
+    def test_with_prefix(self, normalizer):
+        result = normalizer.normalize("Faktura: INV12345")
+        assert result.value is not None
+        assert "INV" in result.value or "12345" in result.value
+
+    def test_year_prefix(self, normalizer):
+        result = normalizer.normalize("2024-12345")
+        assert result.value == "2024-12345"
+        assert result.is_valid is True
+
+    def test_numeric_only(self, normalizer):
+        result = normalizer.normalize("12345678")
+        assert result.value == "12345678"
+        assert result.is_valid is True
+
+    def test_empty_string(self, normalizer):
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_callable(self, normalizer):
+        result = normalizer("A3861")
+        assert result.value == "A3861"
+
+    def test_skip_date_like_sequence(self, normalizer):
+        """Test that 8-digit sequences starting with 20 (dates) are skipped."""
+        result = normalizer.normalize("Invoice 12345 Date 20240115")
+        assert result.value == "12345"
+
+    def test_skip_long_ocr_sequence(self, normalizer):
+        """Test that sequences > 10 digits are skipped."""
+        result = normalizer.normalize("Invoice 54321 OCR 12345678901234")
+        assert result.value == "54321"
+
+    def test_fallback_extraction(self, normalizer):
+        """Test fallback to digit extraction."""
+        # This matches Pattern 3 (short digit sequence 3-10 digits)
+        result = normalizer.normalize("Some text with number 123 embedded")
+        assert result.value == "123"
+        assert result.is_valid is True
+
+    def test_no_valid_sequence(self, normalizer):
+        """Test failure when no valid sequence found."""
+        result = normalizer.normalize("no numbers here")
+        assert result.is_valid is False
+        assert "Cannot extract" in result.error
+
+
+class TestOcrNumberNormalizer:
+    """Tests for OcrNumberNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return OcrNumberNormalizer()
+
+    def test_field_name(self, normalizer):
+        assert normalizer.field_name == "OCR"
+
+    def test_standard_ocr(self, normalizer):
+        result = normalizer.normalize("310196187399952")
+        assert result.value == "310196187399952"
+        assert result.is_valid is True
+
+    def test_with_spaces(self, normalizer):
+        result = normalizer.normalize("3101 9618 7399 952")
+        assert result.value == "310196187399952"
+        assert " " not in result.value
+
+    def test_too_short(self, normalizer):
+        result = normalizer.normalize("1234")
+        assert result.is_valid is False
+
+    def test_empty_string(self, normalizer):
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+
+class TestBankgiroNormalizer:
+    """Tests for BankgiroNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return BankgiroNormalizer()
+
+    def test_field_name(self, normalizer):
+        assert normalizer.field_name == "Bankgiro"
+
+    def test_7_digit_format(self, normalizer):
+        result = normalizer.normalize("782-1713")
+        assert result.value == "782-1713"
+        assert result.is_valid is True
+
+    def test_8_digit_format(self, normalizer):
+        result = normalizer.normalize("5393-9484")
+        assert result.value == "5393-9484"
+        assert result.is_valid is True
+
+    def test_without_dash(self, normalizer):
+        result = normalizer.normalize("7821713")
+        assert result.value is not None
+        assert "-" in result.value
+
+    def test_with_prefix(self, normalizer):
+        result = normalizer.normalize("Bankgiro: 782-1713")
+        assert result.value == "782-1713"
+
+    def test_invalid_too_short(self, normalizer):
+        result = normalizer.normalize("123")
+        assert result.is_valid is False
+
+    def test_empty_string(self, normalizer):
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_invalid_luhn_with_warning(self, normalizer):
+        """Test BG with invalid Luhn checksum returns warning."""
+        # 1234-5679 has invalid Luhn
+        result = normalizer.normalize("1234-5679")
+        assert result.value is not None
+        assert "Luhn checksum failed" in (result.error or "")
+
+    def test_pg_format_excluded(self, normalizer):
+        """Test that PG format (X-X) is not matched as BG."""
+        result = normalizer.normalize("1234567-8")  # PG format
+        assert result.is_valid is False
+
+    def test_raw_7_digits_fallback(self, normalizer):
+        """Test fallback to raw 7 digits without dash."""
+        result = normalizer.normalize("BG number is 7821713 here")
+        assert result.value is not None
+        assert "-" in result.value
+
+    def test_raw_8_digits_invalid_luhn(self, normalizer):
+        """Test raw 8 digits with invalid Luhn."""
+        result = normalizer.normalize("12345679")  # 8 digits, invalid Luhn
+        assert result.value is not None
+        assert "Luhn" in (result.error or "")
+
+
+class TestPlusgiroNormalizer:
+    """Tests for PlusgiroNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return PlusgiroNormalizer()
+
+    def test_field_name(self, normalizer):
+        assert normalizer.field_name == "Plusgiro"
+
+    def test_standard_format(self, normalizer):
+        result = normalizer.normalize("1234567-8")
+        assert result.value is not None
+        assert "-" in result.value
+
+    def test_short_format(self, normalizer):
+        result = normalizer.normalize("12-3")
+        assert result.value is not None
+
+    def test_without_dash(self, normalizer):
+        result = normalizer.normalize("12345678")
+        assert result.value is not None
+        assert "-" in result.value
+
+    def test_with_spaces(self, normalizer):
+        result = normalizer.normalize("486 98 63-6")
+        assert result.value is not None
+
+    def test_empty_string(self, normalizer):
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_invalid_luhn_with_warning(self, normalizer):
+        """Test PG with invalid Luhn returns warning."""
+        result = normalizer.normalize("1234567-9")  # Invalid Luhn
+        assert result.value is not None
+        assert "Luhn checksum failed" in (result.error or "")
+
+    def test_all_digits_fallback(self, normalizer):
+        """Test fallback to all digits extraction."""
+        result = normalizer.normalize("PG 12345")
+        assert result.value is not None
+
+    def test_digit_sequence_fallback(self, normalizer):
+        """Test finding digit sequence in text."""
+        result = normalizer.normalize("Account number: 54321")
+        assert result.value is not None
+
+    def test_too_long_fails(self, normalizer):
+        """Test that > 8 digits fails (no PG format found)."""
+        result = normalizer.normalize("123456789")  # 9 digits, too long
+        # PG is 2-8 digits, so 9 digits is invalid
+        assert result.is_valid is False
+
+    def test_no_digits_fails(self, normalizer):
+        """Test failure when no valid digits found."""
+        result = normalizer.normalize("no numbers")
+        assert result.is_valid is False
+
+    def test_pg_display_format_valid_luhn(self, normalizer):
+        """Test PG display format with valid Luhn checksum."""
+        # 1000009 has valid Luhn checksum
+        result = normalizer.normalize("PG: 100000-9")
+        assert result.value == "100000-9"
+        assert result.is_valid is True
+        assert result.error is None  # No warning for valid Luhn
+
+    def test_pg_all_digits_valid_luhn(self, normalizer):
+        """Test all digits extraction with valid Luhn."""
+        # When no PG format found, extract all digits
+        # 10000008 has valid Luhn (8 digits)
+        result = normalizer.normalize("PG number 10000008")
+        assert result.value == "1000000-8"
+        assert result.is_valid is True
+        assert result.error is None
+
+    def test_pg_digit_sequence_valid_luhn(self, normalizer):
+        """Test digit sequence fallback with valid Luhn."""
+        # Find word-bounded digit sequence
+        # 1000017 has valid Luhn
+        result = normalizer.normalize("Account: 1000017 registered")
+        assert result.value == "100001-7"
+        assert result.is_valid is True
+        assert result.error is None
+
+    def test_pg_digit_sequence_invalid_luhn(self, normalizer):
+        """Test digit sequence fallback with invalid Luhn."""
+        result = normalizer.normalize("Account: 12345678 registered")
+        assert result.value == "1234567-8"
+        assert result.is_valid is True
+        assert "Luhn" in (result.error or "")
+
+    def test_pg_digit_sequence_when_all_digits_too_long(self, normalizer):
+        """Test digit sequence search when all_digits > 8 (lines 79-86)."""
+        # Total digits > 8, so all_digits fallback fails
+        # But there's a word-bounded 7-digit sequence with valid Luhn
+        result = normalizer.normalize("PG is 1000017 but ID is 9999999999")
+        assert result.value == "100001-7"
+        assert result.is_valid is True
+        assert result.error is None  # Valid Luhn
+
+    def test_pg_digit_sequence_invalid_luhn_when_all_digits_too_long(self, normalizer):
+        """Test digit sequence with invalid Luhn when all_digits > 8."""
+        # Total digits > 8, word-bounded sequence has invalid Luhn
+        result = normalizer.normalize("Account 12345 in document 987654321")
+        assert result.value == "1234-5"
+        assert result.is_valid is True
+        assert "Luhn" in (result.error or "")
+
+
+class TestAmountNormalizer:
+    """Tests for AmountNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return AmountNormalizer()
+
+    def test_field_name(self, normalizer):
+        assert normalizer.field_name == "Amount"
+
+    def test_swedish_format(self, normalizer):
+        result = normalizer.normalize("11 699,00")
+        assert result.value is not None
+        assert result.is_valid is True
+
+    def test_with_currency(self, normalizer):
+        result = normalizer.normalize("11 699,00 SEK")
+        assert result.value is not None
+
+    def test_dot_decimal(self, normalizer):
+        result = normalizer.normalize("1234.56")
+        assert result.value == "1234.56"
+
+    def test_integer_amount(self, normalizer):
+        result = normalizer.normalize("Belopp: 11699")
+        assert result.value is not None
+
+    def test_multiple_amounts_returns_last(self, normalizer):
+        result = normalizer.normalize("Subtotal: 100,00\nMoms: 25,00\nTotal: 125,00")
+        assert result.value == "125.00"
+
+    def test_empty_string(self, normalizer):
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_empty_lines_skipped(self, normalizer):
+        """Test that empty lines are skipped."""
+        result = normalizer.normalize("\n\n100,00\n\n")
+        assert result.value == "100.00"
+
+    def test_simple_decimal_fallback(self, normalizer):
+        """Test simple decimal pattern fallback."""
+        result = normalizer.normalize("Price is 99.99 dollars")
+        assert result.value == "99.99"
+
+    def test_standalone_number_fallback(self, normalizer):
+        """Test standalone number >= 3 digits fallback."""
+        result = normalizer.normalize("Amount 12345")
+        assert result.value == "12345.00"
+
+    def test_no_amount_fails(self, normalizer):
+        """Test failure when no amount found."""
+        result = normalizer.normalize("no amount here")
+        assert result.is_valid is False
+
+    def test_value_error_in_amount_parsing(self, normalizer):
+        """Test that ValueError in float conversion is handled."""
+        # A pattern that matches but cannot be converted to float
+        # This is hard to trigger since regex already validates digits
+        result = normalizer.normalize("Amount: abc")
+        assert result.is_valid is False
+
+    def test_shared_validator_fallback(self, normalizer):
+        """Test fallback to shared validator."""
+        # Input that doesn't match primary pattern but shared validator handles
+        result = normalizer.normalize("kr 1234")
+        assert result.value is not None
+
+    def test_simple_decimal_pattern_fallback(self, normalizer):
+        """Test simple decimal pattern fallback."""
+        # Pattern that requires simple_pattern fallback
+        result = normalizer.normalize("Total: 99,99")
+        assert result.value == "99.99"
+
+    def test_integer_pattern_fallback(self, normalizer):
+        """Test integer amount pattern fallback."""
+        result = normalizer.normalize("Amount: 5000")
+        assert result.value == "5000.00"
+
+    def test_standalone_number_fallback(self, normalizer):
+        """Test standalone number >= 3 digits fallback (lines 99-104)."""
+        # No amount/belopp/summa/total keywords, no decimal - reaches standalone pattern
+        result = normalizer.normalize("Reference 12500")
+        assert result.value == "12500.00"
+
+    def test_zero_amount_rejected(self, normalizer):
+        """Test that zero amounts are rejected."""
+        result = normalizer.normalize("0,00 kr")
+        assert result.is_valid is False
+
+    def test_negative_sign_ignored(self, normalizer):
+        """Test that negative sign is ignored (code extracts digits only)."""
+        result = normalizer.normalize("-100,00")
+        # The pattern extracts "100,00" ignoring the negative sign
+        assert result.value == "100.00"
+        assert result.is_valid is True
+
+
+class TestEnhancedAmountNormalizer:
+    """Tests for EnhancedAmountNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return EnhancedAmountNormalizer()
+
+    def test_labeled_amount(self, normalizer):
+        result = normalizer.normalize("Att betala: 1 234,56")
+        assert result.value is not None
+        assert result.is_valid is True
+
+    def test_total_keyword(self, normalizer):
+        result = normalizer.normalize("Total: 9 999,00 kr")
+        assert result.value is not None
+
+    def test_ocr_correction(self, normalizer):
+        # O -> 0 correction
+        result = normalizer.normalize("1O23,45")
+        assert result.value is not None
+
+    def test_summa_keyword(self, normalizer):
+        """Test Swedish 'summa' keyword."""
+        result = normalizer.normalize("Summa: 5 000,00")
+        assert result.value is not None
+
+    def test_moms_lower_priority(self, normalizer):
+        """Test that moms (VAT) has lower priority than summa/total."""
+        # 'summa' keyword has priority 1.0, 'moms' has 0.8
+        result = normalizer.normalize("Moms: 250,00 Summa: 1250,00")
+        assert result.value == "1250.00"
+
+    def test_decimal_pattern_fallback(self, normalizer):
+        """Test decimal pattern extraction."""
+        result = normalizer.normalize("Invoice for 1 234 567,89 kr")
+        assert result.value is not None
+
+    def test_no_amount_fails(self, normalizer):
+        """Test failure when no amount found."""
+        result = normalizer.normalize("no amount")
+        assert result.is_valid is False
+
+    def test_enhanced_empty_string(self, normalizer):
+        """Test empty string fails."""
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_enhanced_shared_validator_fallback(self, normalizer):
+        """Test fallback to shared validator when no labeled patterns match."""
+        # Input that doesn't match labeled patterns but shared validator handles
+        result = normalizer.normalize("kr 1234")
+        assert result.value is not None
+
+    def test_enhanced_decimal_pattern_fallback(self, normalizer):
+        """Test Strategy 4 decimal pattern fallback."""
+        # Input that bypasses labeled patterns and shared validator
+        result = normalizer.normalize("Price: 1 234 567,89")
+        assert result.value is not None
+
+    def test_amount_out_of_range_rejected(self, normalizer):
+        """Test that amounts >= 10,000,000 are rejected."""
+        result = normalizer.normalize("Summa: 99 999 999,00")
+        # Should fail since amount is >= 10,000,000
+        assert result.is_valid is False
+
+    def test_value_error_in_labeled_pattern(self, normalizer):
+        """Test ValueError handling in labeled pattern parsing."""
+        # This is defensive code that's hard to trigger
+        result = normalizer.normalize("Total: abc,00")
+        # Should fall through to other strategies
+        assert result.is_valid is False
+
+    def test_enhanced_decimal_pattern_multiple_amounts(self, normalizer):
+        """Test Strategy 4 with multiple decimal amounts (lines 168-183)."""
+        # Need input that bypasses labeled patterns AND shared validator
+        # but has decimal pattern matches
+        with patch(
+            "inference.pipeline.normalizers.amount.FieldValidators.parse_amount",
+            return_value=None,
+        ):
+            result = normalizer.normalize("Items: 100,00 and 200,00 and 300,00")
+            # Should return max amount
+            assert result.value == "300.00"
+            assert result.is_valid is True
+
+
+class TestDateNormalizer:
+    """Tests for DateNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return DateNormalizer()
+
+    def test_field_name(self, normalizer):
+        assert normalizer.field_name == "Date"
+
+    def test_iso_format(self, normalizer):
+        result = normalizer.normalize("2026-01-31")
+        assert result.value == "2026-01-31"
+        assert result.is_valid is True
+
+    def test_european_dot_format(self, normalizer):
+        result = normalizer.normalize("31.01.2026")
+        assert result.value == "2026-01-31"
+
+    def test_european_slash_format(self, normalizer):
+        result = normalizer.normalize("31/01/2026")
+        assert result.value == "2026-01-31"
+
+    def test_compact_format(self, normalizer):
+        result = normalizer.normalize("20260131")
+        assert result.value == "2026-01-31"
+
+    def test_invalid_date(self, normalizer):
+        result = normalizer.normalize("not a date")
+        assert result.is_valid is False
+
+    def test_empty_string(self, normalizer):
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_dot_format_ymd(self, normalizer):
+        """Test YYYY.MM.DD format."""
+        result = normalizer.normalize("2025.08.29")
+        assert result.value == "2025-08-29"
+
+    def test_invalid_date_value_continues(self, normalizer):
+        """Test that invalid date values are skipped."""
+        result = normalizer.normalize("2025-13-45")  # Invalid month/day
+        assert result.is_valid is False
+
+    def test_year_out_of_range(self, normalizer):
+        """Test that years outside 2000-2100 are rejected."""
+        result = normalizer.normalize("1999-01-01")
+        assert result.is_valid is False
+
+    def test_fallback_pattern_single_digit_day(self, normalizer):
+        """Test fallback pattern with single digit day (European slash format)."""
+        # The shared validator returns None for single digit day like 8/12/2025
+        # So it falls back to the PATTERNS list (European DD/MM/YYYY)
+        result = normalizer.normalize("8/12/2025")
+        assert result.value == "2025-12-08"
+        assert result.is_valid is True
+
+    def test_fallback_pattern_with_mock(self, normalizer):
+        """Test fallback PATTERNS when shared validator returns None (line 83)."""
+        with patch(
+            "inference.pipeline.normalizers.date.FieldValidators.format_date_iso",
+            return_value=None,
+        ):
+            result = normalizer.normalize("2025-08-29")
+            assert result.value == "2025-08-29"
+            assert result.is_valid is True
+
+
+class TestEnhancedDateNormalizer:
+    """Tests for EnhancedDateNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return EnhancedDateNormalizer()
+
+    def test_swedish_text_date(self, normalizer):
+        result = normalizer.normalize("29 december 2024")
+        assert result.value == "2024-12-29"
+        assert result.is_valid is True
+
+    def test_swedish_abbreviated(self, normalizer):
+        result = normalizer.normalize("15 jan 2025")
+        assert result.value == "2025-01-15"
+
+    def test_ocr_correction(self, normalizer):
+        # O -> 0 correction
+        result = normalizer.normalize("2O26-01-31")
+        assert result.value == "2026-01-31"
+
+    def test_empty_string(self, normalizer):
+        """Test empty string fails."""
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_swedish_months(self, normalizer):
+        """Test Swedish month names that work with OCR correction.
+
+        Note: OCRCorrections.correct_digits corrupts some month names:
+        - april -> apr11, juli -> ju11, augusti -> augu571, oktober -> ok706er
+        These months are excluded from this test.
+        """
+        months = [
+            ("15 januari 2025", "2025-01-15"),
+            ("15 februari 2025", "2025-02-15"),
+            ("15 mars 2025", "2025-03-15"),
+            ("15 maj 2025", "2025-05-15"),
+            ("15 juni 2025", "2025-06-15"),
+            ("15 september 2025", "2025-09-15"),
+            ("15 november 2025", "2025-11-15"),
+            ("15 december 2025", "2025-12-15"),
+        ]
+        for text, expected in months:
+            result = normalizer.normalize(text)
+            assert result.value == expected, f"Failed for {text}"
+
+    def test_extended_ymd_slash(self, normalizer):
+        """Test YYYY/MM/DD format."""
+        result = normalizer.normalize("2025/08/29")
+        assert result.value == "2025-08-29"
+
+    def test_extended_dmy_dash(self, normalizer):
+        """Test DD-MM-YYYY format."""
+        result = normalizer.normalize("29-08-2025")
+        assert result.value == "2025-08-29"
+
+    def test_extended_compact(self, normalizer):
+        """Test YYYYMMDD compact format."""
+        result = normalizer.normalize("20250829")
+        assert result.value == "2025-08-29"
+
+    def test_invalid_swedish_month(self, normalizer):
+        """Test invalid Swedish month name falls through."""
+        result = normalizer.normalize("15 invalidmonth 2025")
+        assert result.is_valid is False
+
+    def test_invalid_extended_date_continues(self, normalizer):
+        """Test that invalid dates in extended patterns are skipped."""
+        result = normalizer.normalize("32-13-2025")  # Invalid day/month
+        assert result.is_valid is False
+
+    def test_swedish_pattern_invalid_date(self, normalizer):
+        """Test Swedish pattern with invalid date (Feb 31) falls through.
+
+        When shared validator returns an invalid date like 2025-02-31,
+        is_valid_date returns False, so it tries Swedish pattern,
+        which also fails due to invalid datetime.
+        """
+        result = normalizer.normalize("31 feb 2025")
+        assert result.is_valid is False
+
+    def test_swedish_pattern_year_out_of_range(self, normalizer):
+        """Test Swedish pattern with year outside 2000-2100."""
+        # Use abbreviated month to avoid OCR corruption
+        result = normalizer.normalize("15 jan 1999")
+        # is_valid_date returns False for 1999-01-15, falls through
+        # Swedish pattern matches but year < 2000
+        assert result.is_valid is False
+
+    def test_ymd_compact_format_with_prefix(self, normalizer):
+        """Test YYYYMMDD compact format with surrounding text."""
+        # The compact pattern requires word boundaries
+        result = normalizer.normalize("Date code: 20250315")
+        assert result.value == "2025-03-15"
+
+    def test_swedish_pattern_fallback_with_mock(self, normalizer):
+        """Test Swedish pattern when shared validator returns None (line 170)."""
+        with patch(
+            "inference.pipeline.normalizers.date.FieldValidators.format_date_iso",
+            return_value=None,
+        ):
+            result = normalizer.normalize("15 maj 2025")
+            assert result.value == "2025-05-15"
+            assert result.is_valid is True
+
+    def test_ymd_compact_fallback_with_mock(self, normalizer):
+        """Test ymd_compact pattern when shared validator returns None (lines 187-192)."""
+        with patch(
+            "inference.pipeline.normalizers.date.FieldValidators.format_date_iso",
+            return_value=None,
+        ):
+            result = normalizer.normalize("20250315")
+            assert result.value == "2025-03-15"
+            assert result.is_valid is True
+
+
+class TestSupplierOrgNumberNormalizer:
+    """Tests for SupplierOrgNumberNormalizer."""
+
+    @pytest.fixture
+    def normalizer(self):
+        return SupplierOrgNumberNormalizer()
+
+    def test_field_name(self, normalizer):
+        assert normalizer.field_name == "supplier_org_number"
+
+    def test_standard_format(self, normalizer):
+        result = normalizer.normalize("516406-1102")
+        assert result.value == "516406-1102"
+        assert result.is_valid is True
+
+    def test_with_prefix(self, normalizer):
+        result = normalizer.normalize("Org.nr 516406-1102")
+        assert result.value == "516406-1102"
+
+    def test_without_dash(self, normalizer):
+        result = normalizer.normalize("5164061102")
+        assert result.value == "516406-1102"
+
+    def test_vat_format(self, normalizer):
+        result = normalizer.normalize("SE556123456701")
+        assert result.value is not None
+        assert "-" in result.value
+
+    def test_empty_string(self, normalizer):
+        result = normalizer.normalize("")
+        assert result.is_valid is False
+
+    def test_10_consecutive_digits(self, normalizer):
+        """Test 10 consecutive digits pattern."""
+        result = normalizer.normalize("Company org 5164061102 registered")
+        assert result.value == "516406-1102"
+
+    def test_10_digits_starting_with_zero_accepted(self, normalizer):
+        """Test that 10 digits starting with 0 are accepted by Pattern 1.
+
+        Pattern 1 (NNNNNN-?NNNN) matches any 10 digits with optional dash.
+        Only Pattern 3 (standalone 10 digits) validates first digit != 0.
+        """
+        result = normalizer.normalize("0164061102")
+        assert result.is_valid is True
+        assert result.value == "016406-1102"
+
+    def test_no_org_number_fails(self, normalizer):
+        """Test failure when no org number found."""
+        result = normalizer.normalize("no org number here")
+        assert result.is_valid is False
+
+
+class TestNormalizerRegistry:
+    """Tests for normalizer registry factory."""
+
+    def test_create_registry(self):
+        registry = create_normalizer_registry()
+        assert "InvoiceNumber" in registry
+        assert "OCR" in registry
+        assert "Bankgiro" in registry
+        assert "Plusgiro" in registry
+        assert "Amount" in registry
+        assert "InvoiceDate" in registry
+        assert "InvoiceDueDate" in registry
+        assert "supplier_org_number" in registry
+
+    def test_registry_with_enhanced(self):
+        registry = create_normalizer_registry(use_enhanced=True)
+        # Enhanced normalizers should be used for Amount and Date
+        assert isinstance(registry["Amount"], EnhancedAmountNormalizer)
+        assert isinstance(registry["InvoiceDate"], EnhancedDateNormalizer)
+
+    def test_registry_without_enhanced(self):
+        registry = create_normalizer_registry(use_enhanced=False)
+        assert isinstance(registry["Amount"], AmountNormalizer)
+        assert isinstance(registry["InvoiceDate"], DateNormalizer)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/web/core/__init__.py

@@ -0,0 +1 @@
+"""Tests for web core components."""

tests/web/core/test_task_interface.py

@@ -0,0 +1,672 @@
+"""Tests for unified task management interface.
+
+TDD: These tests are written first (RED phase).
+"""
+
+from abc import ABC
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+class TestTaskStatus:
+    """Tests for TaskStatus dataclass."""
+
+    def test_task_status_basic_fields(self) -> None:
+        """TaskStatus has all required fields."""
+        from inference.web.core.task_interface import TaskStatus
+
+        status = TaskStatus(
+            name="test_runner",
+            is_running=True,
+            pending_count=5,
+            processing_count=2,
+        )
+        assert status.name == "test_runner"
+        assert status.is_running is True
+        assert status.pending_count == 5
+        assert status.processing_count == 2
+
+    def test_task_status_with_error(self) -> None:
+        """TaskStatus can include optional error message."""
+        from inference.web.core.task_interface import TaskStatus
+
+        status = TaskStatus(
+            name="failed_runner",
+            is_running=False,
+            pending_count=0,
+            processing_count=0,
+            error="Connection failed",
+        )
+        assert status.error == "Connection failed"
+
+    def test_task_status_default_error_is_none(self) -> None:
+        """TaskStatus error defaults to None."""
+        from inference.web.core.task_interface import TaskStatus
+
+        status = TaskStatus(
+            name="test",
+            is_running=True,
+            pending_count=0,
+            processing_count=0,
+        )
+        assert status.error is None
+
+    def test_task_status_is_frozen(self) -> None:
+        """TaskStatus is immutable (frozen dataclass)."""
+        from inference.web.core.task_interface import TaskStatus
+
+        status = TaskStatus(
+            name="test",
+            is_running=True,
+            pending_count=0,
+            processing_count=0,
+        )
+        with pytest.raises(AttributeError):
+            status.name = "changed"  # type: ignore[misc]
+
+
+class TestTaskRunnerInterface:
+    """Tests for TaskRunner abstract base class."""
+
+    def test_cannot_instantiate_directly(self) -> None:
+        """TaskRunner is abstract and cannot be instantiated."""
+        from inference.web.core.task_interface import TaskRunner
+
+        with pytest.raises(TypeError):
+            TaskRunner()  # type: ignore[abstract]
+
+    def test_is_abstract_base_class(self) -> None:
+        """TaskRunner inherits from ABC."""
+        from inference.web.core.task_interface import TaskRunner
+
+        assert issubclass(TaskRunner, ABC)
+
+    def test_subclass_missing_name_cannot_instantiate(self) -> None:
+        """Subclass without name property cannot be instantiated."""
+        from inference.web.core.task_interface import TaskRunner, TaskStatus
+
+        class MissingName(TaskRunner):
+            def start(self) -> None:
+                pass
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus("", False, 0, 0)
+
+        with pytest.raises(TypeError):
+            MissingName()  # type: ignore[abstract]
+
+    def test_subclass_missing_start_cannot_instantiate(self) -> None:
+        """Subclass without start method cannot be instantiated."""
+        from inference.web.core.task_interface import TaskRunner, TaskStatus
+
+        class MissingStart(TaskRunner):
+            @property
+            def name(self) -> str:
+                return "test"
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus("", False, 0, 0)
+
+        with pytest.raises(TypeError):
+            MissingStart()  # type: ignore[abstract]
+
+    def test_subclass_missing_stop_cannot_instantiate(self) -> None:
+        """Subclass without stop method cannot be instantiated."""
+        from inference.web.core.task_interface import TaskRunner, TaskStatus
+
+        class MissingStop(TaskRunner):
+            @property
+            def name(self) -> str:
+                return "test"
+
+            def start(self) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus("", False, 0, 0)
+
+        with pytest.raises(TypeError):
+            MissingStop()  # type: ignore[abstract]
+
+    def test_subclass_missing_is_running_cannot_instantiate(self) -> None:
+        """Subclass without is_running property cannot be instantiated."""
+        from inference.web.core.task_interface import TaskRunner, TaskStatus
+
+        class MissingIsRunning(TaskRunner):
+            @property
+            def name(self) -> str:
+                return "test"
+
+            def start(self) -> None:
+                pass
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus("", False, 0, 0)
+
+        with pytest.raises(TypeError):
+            MissingIsRunning()  # type: ignore[abstract]
+
+    def test_subclass_missing_get_status_cannot_instantiate(self) -> None:
+        """Subclass without get_status method cannot be instantiated."""
+        from inference.web.core.task_interface import TaskRunner
+
+        class MissingGetStatus(TaskRunner):
+            @property
+            def name(self) -> str:
+                return "test"
+
+            def start(self) -> None:
+                pass
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+        with pytest.raises(TypeError):
+            MissingGetStatus()  # type: ignore[abstract]
+
+    def test_complete_subclass_can_instantiate(self) -> None:
+        """Complete subclass implementing all methods can be instantiated."""
+        from inference.web.core.task_interface import TaskRunner, TaskStatus
+
+        class CompleteRunner(TaskRunner):
+            def __init__(self) -> None:
+                self._running = False
+
+            @property
+            def name(self) -> str:
+                return "complete_runner"
+
+            def start(self) -> None:
+                self._running = True
+
+            def stop(self, timeout: float | None = None) -> None:
+                self._running = False
+
+            @property
+            def is_running(self) -> bool:
+                return self._running
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus(
+                    name=self.name,
+                    is_running=self._running,
+                    pending_count=0,
+                    processing_count=0,
+                )
+
+        runner = CompleteRunner()
+        assert runner.name == "complete_runner"
+        assert runner.is_running is False
+
+        runner.start()
+        assert runner.is_running is True
+
+        status = runner.get_status()
+        assert status.name == "complete_runner"
+        assert status.is_running is True
+
+        runner.stop()
+        assert runner.is_running is False
+
+
+class TestTaskManager:
+    """Tests for TaskManager facade."""
+
+    def test_register_runner(self) -> None:
+        """Can register a task runner."""
+        from inference.web.core.task_interface import TaskManager, TaskRunner, TaskStatus
+
+        class MockRunner(TaskRunner):
+            @property
+            def name(self) -> str:
+                return "mock"
+
+            def start(self) -> None:
+                pass
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus("mock", False, 0, 0)
+
+        manager = TaskManager()
+        runner = MockRunner()
+        manager.register(runner)
+
+        assert manager.get_runner("mock") is runner
+
+    def test_get_runner_returns_none_for_unknown(self) -> None:
+        """get_runner returns None for unknown runner name."""
+        from inference.web.core.task_interface import TaskManager
+
+        manager = TaskManager()
+        assert manager.get_runner("unknown") is None
+
+    def test_start_all_runners(self) -> None:
+        """start_all starts all registered runners."""
+        from inference.web.core.task_interface import TaskManager, TaskRunner, TaskStatus
+
+        class MockRunner(TaskRunner):
+            def __init__(self, runner_name: str) -> None:
+                self._name = runner_name
+                self._running = False
+
+            @property
+            def name(self) -> str:
+                return self._name
+
+            def start(self) -> None:
+                self._running = True
+
+            def stop(self, timeout: float | None = None) -> None:
+                self._running = False
+
+            @property
+            def is_running(self) -> bool:
+                return self._running
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus(self._name, self._running, 0, 0)
+
+        manager = TaskManager()
+        runner1 = MockRunner("runner1")
+        runner2 = MockRunner("runner2")
+        manager.register(runner1)
+        manager.register(runner2)
+
+        assert runner1.is_running is False
+        assert runner2.is_running is False
+
+        manager.start_all()
+
+        assert runner1.is_running is True
+        assert runner2.is_running is True
+
+    def test_stop_all_runners(self) -> None:
+        """stop_all stops all registered runners."""
+        from inference.web.core.task_interface import TaskManager, TaskRunner, TaskStatus
+
+        class MockRunner(TaskRunner):
+            def __init__(self, runner_name: str) -> None:
+                self._name = runner_name
+                self._running = True
+
+            @property
+            def name(self) -> str:
+                return self._name
+
+            def start(self) -> None:
+                self._running = True
+
+            def stop(self, timeout: float | None = None) -> None:
+                self._running = False
+
+            @property
+            def is_running(self) -> bool:
+                return self._running
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus(self._name, self._running, 0, 0)
+
+        manager = TaskManager()
+        runner1 = MockRunner("runner1")
+        runner2 = MockRunner("runner2")
+        manager.register(runner1)
+        manager.register(runner2)
+
+        assert runner1.is_running is True
+        assert runner2.is_running is True
+
+        manager.stop_all()
+
+        assert runner1.is_running is False
+        assert runner2.is_running is False
+
+    def test_get_all_status(self) -> None:
+        """get_all_status returns status of all runners."""
+        from inference.web.core.task_interface import TaskManager, TaskRunner, TaskStatus
+
+        class MockRunner(TaskRunner):
+            def __init__(self, runner_name: str, pending: int) -> None:
+                self._name = runner_name
+                self._pending = pending
+
+            @property
+            def name(self) -> str:
+                return self._name
+
+            def start(self) -> None:
+                pass
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return True
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus(self._name, True, self._pending, 0)
+
+        manager = TaskManager()
+        manager.register(MockRunner("runner1", 5))
+        manager.register(MockRunner("runner2", 10))
+
+        all_status = manager.get_all_status()
+
+        assert len(all_status) == 2
+        assert all_status["runner1"].pending_count == 5
+        assert all_status["runner2"].pending_count == 10
+
+    def test_get_all_status_empty_when_no_runners(self) -> None:
+        """get_all_status returns empty dict when no runners registered."""
+        from inference.web.core.task_interface import TaskManager
+
+        manager = TaskManager()
+        assert manager.get_all_status() == {}
+
+    def test_runner_names_property(self) -> None:
+        """runner_names returns list of all registered runner names."""
+        from inference.web.core.task_interface import TaskManager, TaskRunner, TaskStatus
+
+        class MockRunner(TaskRunner):
+            def __init__(self, runner_name: str) -> None:
+                self._name = runner_name
+
+            @property
+            def name(self) -> str:
+                return self._name
+
+            def start(self) -> None:
+                pass
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus(self._name, False, 0, 0)
+
+        manager = TaskManager()
+        manager.register(MockRunner("alpha"))
+        manager.register(MockRunner("beta"))
+
+        names = manager.runner_names
+        assert set(names) == {"alpha", "beta"}
+
+    def test_stop_all_with_timeout_distribution(self) -> None:
+        """stop_all distributes timeout across runners."""
+        from inference.web.core.task_interface import TaskManager, TaskRunner, TaskStatus
+
+        received_timeouts: list[float | None] = []
+
+        class MockRunner(TaskRunner):
+            def __init__(self, runner_name: str) -> None:
+                self._name = runner_name
+
+            @property
+            def name(self) -> str:
+                return self._name
+
+            def start(self) -> None:
+                pass
+
+            def stop(self, timeout: float | None = None) -> None:
+                received_timeouts.append(timeout)
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus(self._name, False, 0, 0)
+
+        manager = TaskManager()
+        manager.register(MockRunner("r1"))
+        manager.register(MockRunner("r2"))
+
+        manager.stop_all(timeout=20.0)
+
+        # Timeout should be distributed (20 / 2 = 10 each)
+        assert len(received_timeouts) == 2
+        assert all(t == 10.0 for t in received_timeouts)
+
+    def test_start_all_skips_runners_requiring_arguments(self) -> None:
+        """start_all skips runners that require arguments."""
+        from inference.web.core.task_interface import TaskManager, TaskRunner, TaskStatus
+
+        no_args_started = []
+        with_args_started = []
+
+        class NoArgsRunner(TaskRunner):
+            @property
+            def name(self) -> str:
+                return "no_args"
+
+            def start(self) -> None:
+                no_args_started.append(True)
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus("no_args", False, 0, 0)
+
+        class RequiresArgsRunner(TaskRunner):
+            @property
+            def name(self) -> str:
+                return "requires_args"
+
+            def start(self, handler: object) -> None:  # type: ignore[override]
+                # This runner requires an argument
+                with_args_started.append(True)
+
+            def stop(self, timeout: float | None = None) -> None:
+                pass
+
+            @property
+            def is_running(self) -> bool:
+                return False
+
+            def get_status(self) -> TaskStatus:
+                return TaskStatus("requires_args", False, 0, 0)
+
+        manager = TaskManager()
+        manager.register(NoArgsRunner())
+        manager.register(RequiresArgsRunner())
+
+        # start_all should start no_args runner but skip requires_args
+        manager.start_all()
+
+        assert len(no_args_started) == 1
+        assert len(with_args_started) == 0  # Skipped due to TypeError
+
+    def test_stop_all_with_no_runners(self) -> None:
+        """stop_all does nothing when no runners registered."""
+        from inference.web.core.task_interface import TaskManager
+
+        manager = TaskManager()
+        # Should not raise any exception
+        manager.stop_all()
+        # Just verify it returns without error
+        assert manager.runner_names == []
+
+
+class TestTrainingSchedulerInterface:
+    """Tests for TrainingScheduler implementing TaskRunner."""
+
+    def test_training_scheduler_is_task_runner(self) -> None:
+        """TrainingScheduler inherits from TaskRunner."""
+        from inference.web.core.scheduler import TrainingScheduler
+        from inference.web.core.task_interface import TaskRunner
+
+        scheduler = TrainingScheduler()
+        assert isinstance(scheduler, TaskRunner)
+
+    def test_training_scheduler_name(self) -> None:
+        """TrainingScheduler has correct name."""
+        from inference.web.core.scheduler import TrainingScheduler
+
+        scheduler = TrainingScheduler()
+        assert scheduler.name == "training_scheduler"
+
+    def test_training_scheduler_get_status(self) -> None:
+        """TrainingScheduler provides status via get_status."""
+        from inference.web.core.scheduler import TrainingScheduler
+        from inference.web.core.task_interface import TaskStatus
+
+        scheduler = TrainingScheduler()
+        # Mock the training tasks repository
+        mock_tasks = MagicMock()
+        mock_tasks.get_pending.return_value = [MagicMock(), MagicMock()]
+        scheduler._training_tasks = mock_tasks
+
+        status = scheduler.get_status()
+
+        assert isinstance(status, TaskStatus)
+        assert status.name == "training_scheduler"
+        assert status.is_running is False
+        assert status.pending_count == 2
+
+
+class TestAutoLabelSchedulerInterface:
+    """Tests for AutoLabelScheduler implementing TaskRunner."""
+
+    def test_autolabel_scheduler_is_task_runner(self) -> None:
+        """AutoLabelScheduler inherits from TaskRunner."""
+        from inference.web.core.autolabel_scheduler import AutoLabelScheduler
+        from inference.web.core.task_interface import TaskRunner
+
+        with patch("inference.web.core.autolabel_scheduler.get_storage_helper"):
+            scheduler = AutoLabelScheduler()
+            assert isinstance(scheduler, TaskRunner)
+
+    def test_autolabel_scheduler_name(self) -> None:
+        """AutoLabelScheduler has correct name."""
+        from inference.web.core.autolabel_scheduler import AutoLabelScheduler
+
+        with patch("inference.web.core.autolabel_scheduler.get_storage_helper"):
+            scheduler = AutoLabelScheduler()
+            assert scheduler.name == "autolabel_scheduler"
+
+    def test_autolabel_scheduler_get_status(self) -> None:
+        """AutoLabelScheduler provides status via get_status."""
+        from inference.web.core.autolabel_scheduler import AutoLabelScheduler
+        from inference.web.core.task_interface import TaskStatus
+
+        with patch("inference.web.core.autolabel_scheduler.get_storage_helper"):
+            with patch(
+                "inference.web.core.autolabel_scheduler.get_pending_autolabel_documents"
+            ) as mock_get:
+                mock_get.return_value = [MagicMock(), MagicMock(), MagicMock()]
+
+                scheduler = AutoLabelScheduler()
+                status = scheduler.get_status()
+
+                assert isinstance(status, TaskStatus)
+                assert status.name == "autolabel_scheduler"
+                assert status.is_running is False
+                assert status.pending_count == 3
+
+
+class TestAsyncTaskQueueInterface:
+    """Tests for AsyncTaskQueue implementing TaskRunner."""
+
+    def test_async_queue_is_task_runner(self) -> None:
+        """AsyncTaskQueue inherits from TaskRunner."""
+        from inference.web.workers.async_queue import AsyncTaskQueue
+        from inference.web.core.task_interface import TaskRunner
+
+        queue = AsyncTaskQueue()
+        assert isinstance(queue, TaskRunner)
+
+    def test_async_queue_name(self) -> None:
+        """AsyncTaskQueue has correct name."""
+        from inference.web.workers.async_queue import AsyncTaskQueue
+
+        queue = AsyncTaskQueue()
+        assert queue.name == "async_task_queue"
+
+    def test_async_queue_get_status(self) -> None:
+        """AsyncTaskQueue provides status via get_status."""
+        from inference.web.workers.async_queue import AsyncTaskQueue
+        from inference.web.core.task_interface import TaskStatus
+
+        queue = AsyncTaskQueue()
+        status = queue.get_status()
+
+        assert isinstance(status, TaskStatus)
+        assert status.name == "async_task_queue"
+        assert status.is_running is False
+        assert status.pending_count == 0
+        assert status.processing_count == 0
+
+
+class TestBatchTaskQueueInterface:
+    """Tests for BatchTaskQueue implementing TaskRunner."""
+
+    def test_batch_queue_is_task_runner(self) -> None:
+        """BatchTaskQueue inherits from TaskRunner."""
+        from inference.web.workers.batch_queue import BatchTaskQueue
+        from inference.web.core.task_interface import TaskRunner
+
+        queue = BatchTaskQueue()
+        assert isinstance(queue, TaskRunner)
+
+    def test_batch_queue_name(self) -> None:
+        """BatchTaskQueue has correct name."""
+        from inference.web.workers.batch_queue import BatchTaskQueue
+
+        queue = BatchTaskQueue()
+        assert queue.name == "batch_task_queue"
+
+    def test_batch_queue_get_status(self) -> None:
+        """BatchTaskQueue provides status via get_status."""
+        from inference.web.workers.batch_queue import BatchTaskQueue
+        from inference.web.core.task_interface import TaskStatus
+
+        queue = BatchTaskQueue()
+        status = queue.get_status()
+
+        assert isinstance(status, TaskStatus)
+        assert status.name == "batch_task_queue"
+        assert status.is_running is False
+        assert status.pending_count == 0

tests/web/test_admin_auth.py

@@ -8,80 +8,80 @@ from unittest.mock import MagicMock, patch
 
 from fastapi import HTTPException
 
-from inference.data.admin_db import AdminDB
+from inference.data.repositories import TokenRepository
 from inference.data.admin_models import AdminToken
 from inference.web.core.auth import (
-    get_admin_db,
-    reset_admin_db,
+    get_token_repository,
+    reset_token_repository,
     validate_admin_token,
 )
 
 
 @pytest.fixture
-def mock_admin_db():
-    """Create a mock AdminDB."""
-    db = MagicMock(spec=AdminDB)
-    db.is_valid_admin_token.return_value = True
-    return db
+def mock_token_repo():
+    """Create a mock TokenRepository."""
+    repo = MagicMock(spec=TokenRepository)
+    repo.is_valid.return_value = True
+    return repo
 
 
 @pytest.fixture(autouse=True)
-def reset_db():
-    """Reset admin DB after each test."""
+def reset_repo():
+    """Reset token repository after each test."""
     yield
-    reset_admin_db()
+    reset_token_repository()
 
 
 class TestValidateAdminToken:
     """Tests for validate_admin_token dependency."""
 
-    def test_missing_token_raises_401(self, mock_admin_db):
+    def test_missing_token_raises_401(self, mock_token_repo):
         """Test that missing token raises 401."""
         import asyncio
 
         with pytest.raises(HTTPException) as exc_info:
             asyncio.get_event_loop().run_until_complete(
-                validate_admin_token(None, mock_admin_db)
+                validate_admin_token(None, mock_token_repo)
             )
 
         assert exc_info.value.status_code == 401
         assert "Admin token required" in exc_info.value.detail
 
-    def test_invalid_token_raises_401(self, mock_admin_db):
+    def test_invalid_token_raises_401(self, mock_token_repo):
         """Test that invalid token raises 401."""
         import asyncio
 
-        mock_admin_db.is_valid_admin_token.return_value = False
+        mock_token_repo.is_valid.return_value = False
 
         with pytest.raises(HTTPException) as exc_info:
             asyncio.get_event_loop().run_until_complete(
-                validate_admin_token("invalid-token", mock_admin_db)
+                validate_admin_token("invalid-token", mock_token_repo)
             )
 
         assert exc_info.value.status_code == 401
         assert "Invalid or expired" in exc_info.value.detail
 
-    def test_valid_token_returns_token(self, mock_admin_db):
+    def test_valid_token_returns_token(self, mock_token_repo):
         """Test that valid token is returned."""
         import asyncio
 
         token = "valid-test-token"
-        mock_admin_db.is_valid_admin_token.return_value = True
+        mock_token_repo.is_valid.return_value = True
 
         result = asyncio.get_event_loop().run_until_complete(
-            validate_admin_token(token, mock_admin_db)
+            validate_admin_token(token, mock_token_repo)
         )
 
         assert result == token
-        mock_admin_db.update_admin_token_usage.assert_called_once_with(token)
+        mock_token_repo.update_usage.assert_called_once_with(token)
 
 
-class TestAdminDB:
-    """Tests for AdminDB operations."""
+class TestTokenRepository:
+    """Tests for TokenRepository operations."""
 
-    def test_is_valid_admin_token_active(self):
+    def test_is_valid_active_token(self):
         """Test valid active token."""
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.token_repository.BaseRepository._session") as mock_ctx:
             mock_session = MagicMock()
             mock_ctx.return_value.__enter__.return_value = mock_session
 
@@ -93,12 +93,12 @@ class TestAdminDB:
             )
             mock_session.get.return_value = mock_token
 
-            db = AdminDB()
-            assert db.is_valid_admin_token("test-token") is True
+            repo = TokenRepository()
+            assert repo.is_valid("test-token") is True
 
-    def test_is_valid_admin_token_inactive(self):
+    def test_is_valid_inactive_token(self):
         """Test inactive token."""
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.token_repository.BaseRepository._session") as mock_ctx:
             mock_session = MagicMock()
             mock_ctx.return_value.__enter__.return_value = mock_session
 
@@ -110,12 +110,12 @@ class TestAdminDB:
             )
             mock_session.get.return_value = mock_token
 
-            db = AdminDB()
-            assert db.is_valid_admin_token("test-token") is False
+            repo = TokenRepository()
+            assert repo.is_valid("test-token") is False
 
-    def test_is_valid_admin_token_expired(self):
+    def test_is_valid_expired_token(self):
         """Test expired token."""
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.token_repository.BaseRepository._session") as mock_ctx:
             mock_session = MagicMock()
             mock_ctx.return_value.__enter__.return_value = mock_session
 
@@ -127,36 +127,38 @@ class TestAdminDB:
             )
             mock_session.get.return_value = mock_token
 
-            db = AdminDB()
-            assert db.is_valid_admin_token("test-token") is False
+            repo = TokenRepository()
+            # Need to also mock _now() to ensure proper comparison
+            with patch.object(repo, "_now", return_value=datetime.utcnow()):
+                assert repo.is_valid("test-token") is False
 
-    def test_is_valid_admin_token_not_found(self):
+    def test_is_valid_token_not_found(self):
         """Test token not found."""
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.token_repository.BaseRepository._session") as mock_ctx:
             mock_session = MagicMock()
             mock_ctx.return_value.__enter__.return_value = mock_session
             mock_session.get.return_value = None
 
-            db = AdminDB()
-            assert db.is_valid_admin_token("nonexistent") is False
+            repo = TokenRepository()
+            assert repo.is_valid("nonexistent") is False
 
 
-class TestGetAdminDb:
-    """Tests for get_admin_db function."""
+class TestGetTokenRepository:
+    """Tests for get_token_repository function."""
 
     def test_returns_singleton(self):
-        """Test that get_admin_db returns singleton."""
-        reset_admin_db()
+        """Test that get_token_repository returns singleton."""
+        reset_token_repository()
 
-        db1 = get_admin_db()
-        db2 = get_admin_db()
+        repo1 = get_token_repository()
+        repo2 = get_token_repository()
 
-        assert db1 is db2
+        assert repo1 is repo2
 
     def test_reset_clears_singleton(self):
         """Test that reset clears singleton."""
-        db1 = get_admin_db()
-        reset_admin_db()
-        db2 = get_admin_db()
+        repo1 = get_token_repository()
+        reset_token_repository()
+        repo2 = get_token_repository()
 
-        assert db1 is not db2
+        assert repo1 is not repo2

tests/web/test_admin_routes_enhanced.py

@@ -11,7 +11,12 @@ from fastapi.testclient import TestClient
 
 from inference.web.api.v1.admin.documents import create_documents_router
 from inference.web.config import StorageConfig
-from inference.web.core.auth import validate_admin_token, get_admin_db
+from inference.web.core.auth import (
+    validate_admin_token,
+    get_document_repository,
+    get_annotation_repository,
+    get_training_task_repository,
+)
 
 
 class MockAdminDocument:
@@ -59,14 +64,14 @@ class MockAnnotation:
         self.created_at = kwargs.get('created_at', datetime.utcnow())
 
 
-class MockAdminDB:
-    """Mock AdminDB for testing enhanced features."""
+class MockDocumentRepository:
+    """Mock DocumentRepository for testing enhanced features."""
 
     def __init__(self):
         self.documents = {}
-        self.annotations = {}
+        self.annotations = {}  # Shared reference for filtering
 
-    def get_documents_by_token(
+    def get_paginated(
         self,
         admin_token=None,
         status=None,
@@ -103,32 +108,51 @@ class MockAdminDB:
         total = len(docs)
         return docs[offset:offset+limit], total
 
-    def get_annotations_for_document(self, document_id):
-        """Get annotations for document."""
-        return self.annotations.get(str(document_id), [])
-
-    def count_documents_by_status(self, admin_token):
+    def count_by_status(self, admin_token=None):
         """Count documents by status."""
         counts = {}
         for doc in self.documents.values():
-            if doc.admin_token == admin_token:
+            if admin_token is None or doc.admin_token == admin_token:
                 counts[doc.status] = counts.get(doc.status, 0) + 1
         return counts
 
-    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
 
+
+class MockAnnotationRepository:
+    """Mock AnnotationRepository for testing enhanced features."""
+
+    def __init__(self):
+        self.annotations = {}
+
+    def get_for_document(self, document_id, page_number=None):
+        """Get annotations for document."""
+        return self.annotations.get(str(document_id), [])
+
+
+class MockTrainingTaskRepository:
+    """Mock TrainingTaskRepository for testing enhanced features."""
+
+    def __init__(self):
+        self.training_tasks = {}
+        self.training_links = {}
+
     def get_document_training_tasks(self, document_id):
         """Get training tasks that used this document."""
-        return []  # No training history in this test
+        return self.training_links.get(str(document_id), [])
 
-    def get_training_task(self, task_id):
+    def get(self, task_id):
         """Get training task by ID."""
-        return None  # No training tasks in this test
+        return self.training_tasks.get(str(task_id))
 
 
 @pytest.fixture
@@ -136,8 +160,10 @@ 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()
+    mock_training_task_repo = MockTrainingTaskRepository()
 
     # Add test documents
     doc1 = MockAdminDocument(
@@ -162,19 +188,19 @@ def app():
         batch_id=None
     )
 
-    mock_db.documents[str(doc1.document_id)] = doc1
-    mock_db.documents[str(doc2.document_id)] = doc2
-    mock_db.documents[str(doc3.document_id)] = doc3
+    mock_document_repo.documents[str(doc1.document_id)] = doc1
+    mock_document_repo.documents[str(doc2.document_id)] = doc2
+    mock_document_repo.documents[str(doc3.document_id)] = doc3
 
     # Add annotations to doc1 and doc2
-    mock_db.annotations[str(doc1.document_id)] = [
+    mock_annotation_repo.annotations[str(doc1.document_id)] = [
         MockAnnotation(
             document_id=doc1.document_id,
             class_name="invoice_number",
             text_value="INV-001"
         )
     ]
-    mock_db.annotations[str(doc2.document_id)] = [
+    mock_annotation_repo.annotations[str(doc2.document_id)] = [
         MockAnnotation(
             document_id=doc2.document_id,
             class_id=6,
@@ -189,9 +215,14 @@ def app():
         )
     ]
 
+    # Share annotation data with document repo for filtering
+    mock_document_repo.annotations = mock_annotation_repo.annotations
+
     # 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
+    app.dependency_overrides[get_annotation_repository] = lambda: mock_annotation_repo
+    app.dependency_overrides[get_training_task_repository] = lambda: mock_training_task_repo
 
     # Include router
     router = create_documents_router(StorageConfig())

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:

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()

tests/web/test_augmentation_routes.py

@@ -11,7 +11,11 @@ from fastapi.testclient import TestClient
 import numpy as np
 
 from inference.web.api.v1.admin.augmentation import create_augmentation_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,
+    get_dataset_repository,
+)
 
 
 TEST_ADMIN_TOKEN = "test-admin-token-12345"
@@ -26,18 +30,27 @@ def admin_token() -> str:
 
 
 @pytest.fixture
-def mock_admin_db() -> MagicMock:
-    """Create a mock AdminDB for testing."""
+def mock_document_repo() -> MagicMock:
+    """Create a mock DocumentRepository for testing."""
     mock = MagicMock()
     # Default return values
-    mock.get_document_by_token.return_value = None
-    mock.get_dataset.return_value = None
-    mock.get_augmented_datasets.return_value = ([], 0)
+    mock.get.return_value = None
+    mock.get_by_token.return_value = None
     return mock
 
 
 @pytest.fixture
-def admin_client(mock_admin_db: MagicMock) -> TestClient:
+def mock_dataset_repo() -> MagicMock:
+    """Create a mock DatasetRepository for testing."""
+    mock = MagicMock()
+    # Default return values
+    mock.get.return_value = None
+    mock.get_paginated.return_value = ([], 0)
+    return mock
+
+
+@pytest.fixture
+def admin_client(mock_document_repo: MagicMock, mock_dataset_repo: MagicMock) -> TestClient:
     """Create test client with admin authentication."""
     app = FastAPI()
 
@@ -45,11 +58,15 @@ def admin_client(mock_admin_db: MagicMock) -> TestClient:
     def get_token_override():
         return TEST_ADMIN_TOKEN
 
-    def get_db_override():
-        return mock_admin_db
+    def get_document_repo_override():
+        return mock_document_repo
+
+    def get_dataset_repo_override():
+        return mock_dataset_repo
 
     app.dependency_overrides[validate_admin_token] = get_token_override
-    app.dependency_overrides[get_admin_db] = get_db_override
+    app.dependency_overrides[get_document_repository] = get_document_repo_override
+    app.dependency_overrides[get_dataset_repository] = get_dataset_repo_override
 
     # Include router - the router already has /augmentation prefix
     # so we add /api/v1/admin to get /api/v1/admin/augmentation
@@ -60,15 +77,19 @@ def admin_client(mock_admin_db: MagicMock) -> TestClient:
 
 
 @pytest.fixture
-def unauthenticated_client(mock_admin_db: MagicMock) -> TestClient:
+def unauthenticated_client(mock_document_repo: MagicMock, mock_dataset_repo: MagicMock) -> TestClient:
     """Create test client WITHOUT admin authentication override."""
     app = FastAPI()
 
-    # Only override the database, NOT the token validation
-    def get_db_override():
-        return mock_admin_db
+    # Only override the repositories, NOT the token validation
+    def get_document_repo_override():
+        return mock_document_repo
 
-    app.dependency_overrides[get_admin_db] = get_db_override
+    def get_dataset_repo_override():
+        return mock_dataset_repo
+
+    app.dependency_overrides[get_document_repository] = get_document_repo_override
+    app.dependency_overrides[get_dataset_repository] = get_dataset_repo_override
 
     router = create_augmentation_router()
     app.include_router(router, prefix="/api/v1/admin")
@@ -142,13 +163,13 @@ class TestAugmentationPreviewEndpoint:
         admin_client: TestClient,
         admin_token: str,
         sample_document_id: str,
-        mock_admin_db: MagicMock,
+        mock_document_repo: MagicMock,
     ) -> None:
         """Test previewing augmentation on a document."""
         # Mock document exists
         mock_document = MagicMock()
         mock_document.images_dir = "/fake/path"
-        mock_admin_db.get_document.return_value = mock_document
+        mock_document_repo.get.return_value = mock_document
 
         # Create a fake image (100x100 RGB)
         fake_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
@@ -218,13 +239,13 @@ class TestAugmentationPreviewConfigEndpoint:
         admin_client: TestClient,
         admin_token: str,
         sample_document_id: str,
-        mock_admin_db: MagicMock,
+        mock_document_repo: MagicMock,
     ) -> None:
         """Test previewing full config on a document."""
         # Mock document exists
         mock_document = MagicMock()
         mock_document.images_dir = "/fake/path"
-        mock_admin_db.get_document.return_value = mock_document
+        mock_document_repo.get.return_value = mock_document
 
         # Create a fake image (100x100 RGB)
         fake_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
@@ -260,13 +281,13 @@ class TestAugmentationBatchEndpoint:
         admin_client: TestClient,
         admin_token: str,
         sample_dataset_id: str,
-        mock_admin_db: MagicMock,
+        mock_dataset_repo: MagicMock,
     ) -> None:
         """Test creating augmented dataset."""
         # Mock dataset exists
         mock_dataset = MagicMock()
         mock_dataset.total_images = 100
-        mock_admin_db.get_dataset.return_value = mock_dataset
+        mock_dataset_repo.get.return_value = mock_dataset
 
         response = admin_client.post(
             "/api/v1/admin/augmentation/batch",

tests/web/test_autolabel_with_locks.py

@@ -9,7 +9,6 @@ from unittest.mock import Mock, MagicMock
 from uuid import uuid4
 
 from inference.web.services.autolabel import AutoLabelService
-from inference.data.admin_db import AdminDB
 
 
 class MockDocument:
@@ -23,19 +22,18 @@ class MockDocument:
         self.auto_label_error = None
 
 
-class MockAdminDB:
-    """Mock AdminDB for testing."""
+class MockDocumentRepository:
+    """Mock DocumentRepository for testing."""
 
     def __init__(self):
         self.documents = {}
-        self.annotations = []
         self.status_updates = []
 
-    def get_document(self, document_id):
+    def get(self, document_id):
         """Get document by ID."""
         return self.documents.get(str(document_id))
 
-    def update_document_status(
+    def update_status(
         self,
         document_id,
         status=None,
@@ -58,19 +56,32 @@ class MockAdminDB:
             if auto_label_error:
                 doc.auto_label_error = auto_label_error
 
-    def delete_annotations_for_document(self, document_id, source=None):
+
+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_annotations_batch(self, annotations):
+    def create_batch(self, annotations):
         """Mock create annotations."""
         self.annotations.extend(annotations)
 
 
 @pytest.fixture
-def mock_db():
-    """Create mock admin DB."""
-    return MockAdminDB()
+def mock_doc_repo():
+    """Create mock document repository."""
+    return MockDocumentRepository()
+
+
+@pytest.fixture
+def mock_ann_repo():
+    """Create mock annotation repository."""
+    return MockAnnotationRepository()
 
 
 @pytest.fixture
@@ -82,10 +93,14 @@ def auto_label_service(monkeypatch):
     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):
+    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
 
@@ -93,11 +108,11 @@ def auto_label_service(monkeypatch):
 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):
+    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_db.documents[document_id] = MockDocument(
+        mock_doc_repo.documents[document_id] = MockDocument(
             document_id=document_id,
             annotation_lock_until=None,
         )
@@ -111,21 +126,22 @@ class TestAutoLabelWithLocks:
             document_id=document_id,
             file_path=str(test_file),
             field_values={"invoice_number": "INV-001"},
-            db=mock_db,
+            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_db.status_updates) >= 2
-        assert mock_db.status_updates[0]["auto_label_status"] == "running"
+        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_db, tmp_path):
+    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_db.documents[document_id] = MockDocument(
+        mock_doc_repo.documents[document_id] = MockDocument(
             document_id=document_id,
             annotation_lock_until=lock_until,
         )
@@ -139,7 +155,8 @@ class TestAutoLabelWithLocks:
             document_id=document_id,
             file_path=str(test_file),
             field_values={"invoice_number": "INV-001"},
-            db=mock_db,
+            doc_repo=mock_doc_repo,
+            ann_repo=mock_ann_repo,
         )
 
         # Should fail
@@ -150,15 +167,15 @@ class TestAutoLabelWithLocks:
         # Verify status was updated to failed
         assert any(
             update["auto_label_status"] == "failed"
-            for update in mock_db.status_updates
+            for update in mock_doc_repo.status_updates
         )
 
-    def test_auto_label_expired_lock_succeeds(self, auto_label_service, mock_db, tmp_path):
+    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_db.documents[document_id] = MockDocument(
+        mock_doc_repo.documents[document_id] = MockDocument(
             document_id=document_id,
             annotation_lock_until=lock_until,
         )
@@ -172,18 +189,19 @@ class TestAutoLabelWithLocks:
             document_id=document_id,
             file_path=str(test_file),
             field_values={"invoice_number": "INV-001"},
-            db=mock_db,
+            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_db, tmp_path):
+    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_db.documents[document_id] = MockDocument(
+        mock_doc_repo.documents[document_id] = MockDocument(
             document_id=document_id,
             annotation_lock_until=lock_until,
         )
@@ -197,14 +215,15 @@ class TestAutoLabelWithLocks:
             document_id=document_id,
             file_path=str(test_file),
             field_values={"invoice_number": "INV-001"},
-            db=mock_db,
+            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_db, tmp_path):
+    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"
@@ -215,19 +234,20 @@ class TestAutoLabelWithLocks:
             document_id=str(uuid4()),
             file_path=str(test_file),
             field_values={"invoice_number": "INV-001"},
-            db=mock_db,
+            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_db, tmp_path):
+    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_db.documents[document_id] = MockDocument(
+        mock_doc_repo.documents[document_id] = MockDocument(
             document_id=document_id,
             annotation_lock_until=lock_until,
         )
@@ -241,7 +261,8 @@ class TestAutoLabelWithLocks:
             document_id=document_id,
             file_path=str(test_file),
             field_values={"invoice_number": "INV-001"},
-            db=mock_db,
+            doc_repo=mock_doc_repo,
+            ann_repo=mock_ann_repo,
             # skip_lock_check not specified, should default to False
         )
 

tests/web/test_batch_upload_routes.py

@@ -11,20 +11,20 @@ import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
-from inference.web.api.v1.batch.routes import router
-from inference.web.core.auth import validate_admin_token, get_admin_db
+from inference.web.api.v1.batch.routes import router, get_batch_repository
+from inference.web.core.auth import validate_admin_token
 from inference.web.workers.batch_queue import init_batch_queue, shutdown_batch_queue
 from inference.web.services.batch_upload import BatchUploadService
 
 
-class MockAdminDB:
-    """Mock AdminDB for testing."""
+class MockBatchUploadRepository:
+    """Mock BatchUploadRepository for testing."""
 
     def __init__(self):
         self.batches = {}
         self.batch_files = {}
 
-    def create_batch_upload(self, admin_token, filename, file_size, upload_source):
+    def create(self, admin_token, filename, file_size, upload_source="ui"):
         batch_id = uuid4()
         batch = type('BatchUpload', (), {
             'batch_id': batch_id,
@@ -46,13 +46,13 @@ class MockAdminDB:
         self.batches[batch_id] = batch
         return batch
 
-    def update_batch_upload(self, batch_id, **kwargs):
+    def update(self, batch_id, **kwargs):
         if batch_id in self.batches:
             batch = self.batches[batch_id]
             for key, value in kwargs.items():
                 setattr(batch, key, value)
 
-    def create_batch_upload_file(self, batch_id, filename, **kwargs):
+    def create_file(self, batch_id, filename, **kwargs):
         file_id = uuid4()
         defaults = {
             'file_id': file_id,
@@ -70,7 +70,7 @@ class MockAdminDB:
         self.batch_files[batch_id].append(file_record)
         return file_record
 
-    def update_batch_upload_file(self, file_id, **kwargs):
+    def update_file(self, file_id, **kwargs):
         for files in self.batch_files.values():
             for file_record in files:
                 if file_record.file_id == file_id:
@@ -78,7 +78,7 @@ class MockAdminDB:
                         setattr(file_record, key, value)
                     return
 
-    def get_batch_upload(self, batch_id):
+    def get(self, batch_id):
         return self.batches.get(batch_id, type('BatchUpload', (), {
             'batch_id': batch_id,
             'admin_token': 'test-token',
@@ -95,12 +95,15 @@ class MockAdminDB:
             'completed_at': datetime.utcnow(),
         })())
 
-    def get_batch_upload_files(self, batch_id):
+    def get_files(self, batch_id):
         return self.batch_files.get(batch_id, [])
 
-    def get_batch_uploads_by_token(self, admin_token, limit=50, offset=0):
+    def get_paginated(self, admin_token=None, limit=50, offset=0):
         """Get batches filtered by admin token with pagination."""
-        token_batches = [b for b in self.batches.values() if b.admin_token == admin_token]
+        if admin_token:
+            token_batches = [b for b in self.batches.values() if b.admin_token == admin_token]
+        else:
+            token_batches = list(self.batches.values())
         total = len(token_batches)
         return token_batches[offset:offset+limit], total
 
@@ -110,15 +113,15 @@ def app():
     """Create test FastAPI app with mocked dependencies."""
     app = FastAPI()
 
-    # Create mock admin DB
-    mock_admin_db = MockAdminDB()
+    # Create mock batch upload repository
+    mock_batch_upload_repo = MockBatchUploadRepository()
 
     # Override dependencies
     app.dependency_overrides[validate_admin_token] = lambda: "test-token"
-    app.dependency_overrides[get_admin_db] = lambda: mock_admin_db
+    app.dependency_overrides[get_batch_repository] = lambda: mock_batch_upload_repo
 
     # Initialize batch queue with mock service
-    batch_service = BatchUploadService(mock_admin_db)
+    batch_service = BatchUploadService(mock_batch_upload_repo)
     init_batch_queue(batch_service)
 
     app.include_router(router)

tests/web/test_batch_upload_service.py

@@ -9,19 +9,18 @@ from uuid import uuid4
 
 import pytest
 
-from inference.data.admin_db import AdminDB
 from inference.web.services.batch_upload import BatchUploadService
 
 
 @pytest.fixture
-def admin_db():
-    """Mock admin database for testing."""
-    class MockAdminDB:
+def batch_repo():
+    """Mock batch upload repository for testing."""
+    class MockBatchUploadRepository:
         def __init__(self):
             self.batches = {}
             self.batch_files = {}
 
-        def create_batch_upload(self, admin_token, filename, file_size, upload_source):
+        def create(self, admin_token, filename, file_size, upload_source):
             batch_id = uuid4()
             batch = type('BatchUpload', (), {
                 'batch_id': batch_id,
@@ -43,13 +42,13 @@ def admin_db():
             self.batches[batch_id] = batch
             return batch
 
-        def update_batch_upload(self, batch_id, **kwargs):
+        def update(self, batch_id, **kwargs):
             if batch_id in self.batches:
                 batch = self.batches[batch_id]
                 for key, value in kwargs.items():
                     setattr(batch, key, value)
 
-        def create_batch_upload_file(self, batch_id, filename, **kwargs):
+        def create_file(self, batch_id, filename, **kwargs):
             file_id = uuid4()
             # Set defaults for attributes
             defaults = {
@@ -68,7 +67,7 @@ def admin_db():
             self.batch_files[batch_id].append(file_record)
             return file_record
 
-        def update_batch_upload_file(self, file_id, **kwargs):
+        def update_file(self, file_id, **kwargs):
             for files in self.batch_files.values():
                 for file_record in files:
                     if file_record.file_id == file_id:
@@ -76,19 +75,19 @@ def admin_db():
                             setattr(file_record, key, value)
                         return
 
-        def get_batch_upload(self, batch_id):
+        def get(self, batch_id):
             return self.batches.get(batch_id)
 
-        def get_batch_upload_files(self, batch_id):
+        def get_files(self, batch_id):
             return self.batch_files.get(batch_id, [])
 
-    return MockAdminDB()
+    return MockBatchUploadRepository()
 
 
 @pytest.fixture
-def batch_service(admin_db):
+def batch_service(batch_repo):
     """Batch upload service instance."""
-    return BatchUploadService(admin_db)
+    return BatchUploadService(batch_repo)
 
 
 def create_test_zip(files):
@@ -194,7 +193,7 @@ INV002,F2024-002,2024-01-16,2500.00,7350087654321,123-4567,C124
         assert csv_data["INV001"]["Amount"] == "1500.00"
         assert csv_data["INV001"]["customer_number"] == "C123"
 
-    def test_get_batch_status(self, batch_service, admin_db):
+    def test_get_batch_status(self, batch_service, batch_repo):
         """Test getting batch upload status."""
         # Create a batch
         zip_content = create_test_zip({"INV001.pdf": b"%PDF-1.4 test"})

tests/web/test_dataset_builder.py

@@ -16,7 +16,6 @@ from inference.data.admin_models import (
     AdminAnnotation,
     AdminDocument,
     TrainingDataset,
-    FIELD_CLASSES,
 )
 
 
@@ -35,10 +34,10 @@ def tmp_admin_images(tmp_path):
 
 
 @pytest.fixture
-def mock_admin_db():
-    """Mock AdminDB with dataset and document methods."""
-    db = MagicMock()
-    db.create_dataset.return_value = TrainingDataset(
+def mock_datasets_repo():
+    """Mock DatasetRepository."""
+    repo = MagicMock()
+    repo.create.return_value = TrainingDataset(
         dataset_id=uuid4(),
         name="test-dataset",
         status="building",
@@ -46,7 +45,19 @@ def mock_admin_db():
         val_ratio=0.1,
         seed=42,
     )
-    return db
+    return repo
+
+
+@pytest.fixture
+def mock_documents_repo():
+    """Mock DocumentRepository."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_annotations_repo():
+    """Mock AnnotationRepository."""
+    return MagicMock()
 
 
 @pytest.fixture
@@ -60,6 +71,7 @@ def sample_documents(tmp_admin_images):
         doc.filename = f"{doc_id}.pdf"
         doc.page_count = 2
         doc.file_path = str(tmp_path / "admin_images" / str(doc_id))
+        doc.group_key = None  # Default to no group
         docs.append(doc)
     return docs
 
@@ -89,21 +101,27 @@ class TestDatasetBuilder:
     """Tests for DatasetBuilder."""
 
     def test_build_creates_directory_structure(
-        self, tmp_path, mock_admin_db, sample_documents, sample_annotations
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo,
+        sample_documents, sample_annotations
     ):
         """Dataset builder should create images/ and labels/ with train/val/test subdirs."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
         dataset_dir = tmp_path / "datasets" / "test"
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
-        # Mock DB calls
-        mock_admin_db.get_documents_by_ids.return_value = sample_documents
-        mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+        # Mock repo calls
+        mock_documents_repo.get_by_ids.return_value = sample_documents
+        mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
             sample_annotations.get(str(doc_id), [])
         )
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in sample_documents],
@@ -119,18 +137,24 @@ class TestDatasetBuilder:
             assert (result_dir / "labels" / split).exists()
 
     def test_build_copies_images(
-        self, tmp_path, mock_admin_db, sample_documents, sample_annotations
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo,
+        sample_documents, sample_annotations
     ):
         """Images should be copied from admin_images to dataset folder."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = sample_documents
-        mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = sample_documents
+        mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
             sample_annotations.get(str(doc_id), [])
         )
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         result = builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in sample_documents],
@@ -149,18 +173,24 @@ class TestDatasetBuilder:
         assert total_images == 10  # 5 docs * 2 pages
 
     def test_build_generates_yolo_labels(
-        self, tmp_path, mock_admin_db, sample_documents, sample_annotations
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo,
+        sample_documents, sample_annotations
     ):
         """YOLO label files should be generated with correct format."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = sample_documents
-        mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = sample_documents
+        mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
             sample_annotations.get(str(doc_id), [])
         )
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in sample_documents],
@@ -187,18 +217,24 @@ class TestDatasetBuilder:
         assert 0 <= float(parts[2]) <= 1  # y_center
 
     def test_build_generates_data_yaml(
-        self, tmp_path, mock_admin_db, sample_documents, sample_annotations
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo,
+        sample_documents, sample_annotations
     ):
         """data.yaml should be generated with correct field classes."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = sample_documents
-        mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = sample_documents
+        mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
             sample_annotations.get(str(doc_id), [])
         )
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in sample_documents],
@@ -217,18 +253,24 @@ class TestDatasetBuilder:
         assert "invoice_number" in content
 
     def test_build_splits_documents_correctly(
-        self, tmp_path, mock_admin_db, sample_documents, sample_annotations
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo,
+        sample_documents, sample_annotations
     ):
         """Documents should be split into train/val/test according to ratios."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = sample_documents
-        mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = sample_documents
+        mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
             sample_annotations.get(str(doc_id), [])
         )
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in sample_documents],
@@ -238,8 +280,8 @@ class TestDatasetBuilder:
             admin_images_dir=tmp_path / "admin_images",
         )
 
-        # Verify add_dataset_documents was called with correct splits
-        call_args = mock_admin_db.add_dataset_documents.call_args
+        # Verify add_documents was called with correct splits
+        call_args = mock_datasets_repo.add_documents.call_args
         docs_added = call_args[1]["documents"] if "documents" in call_args[1] else call_args[0][1]
         splits = [d["split"] for d in docs_added]
         assert "train" in splits
@@ -248,18 +290,24 @@ class TestDatasetBuilder:
         assert train_count >= 3  # At least 3 of 5 should be train
 
     def test_build_updates_status_to_ready(
-        self, tmp_path, mock_admin_db, sample_documents, sample_annotations
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo,
+        sample_documents, sample_annotations
     ):
         """After successful build, dataset status should be updated to 'ready'."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = sample_documents
-        mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = sample_documents
+        mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
             sample_annotations.get(str(doc_id), [])
         )
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in sample_documents],
@@ -269,22 +317,27 @@ class TestDatasetBuilder:
             admin_images_dir=tmp_path / "admin_images",
         )
 
-        mock_admin_db.update_dataset_status.assert_called_once()
-        call_kwargs = mock_admin_db.update_dataset_status.call_args[1]
+        mock_datasets_repo.update_status.assert_called_once()
+        call_kwargs = mock_datasets_repo.update_status.call_args[1]
         assert call_kwargs["status"] == "ready"
         assert call_kwargs["total_documents"] == 5
         assert call_kwargs["total_images"] == 10
 
     def test_build_sets_failed_on_error(
-        self, tmp_path, mock_admin_db
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
     ):
         """If build fails, dataset status should be set to 'failed'."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = []  # No docs found
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = []  # No docs found
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         with pytest.raises(ValueError):
             builder.build_dataset(
                 dataset_id=str(dataset.dataset_id),
@@ -295,27 +348,33 @@ class TestDatasetBuilder:
                 admin_images_dir=tmp_path / "admin_images",
             )
 
-        mock_admin_db.update_dataset_status.assert_called_once()
-        call_kwargs = mock_admin_db.update_dataset_status.call_args[1]
+        mock_datasets_repo.update_status.assert_called_once()
+        call_kwargs = mock_datasets_repo.update_status.call_args[1]
         assert call_kwargs["status"] == "failed"
 
     def test_build_with_seed_produces_deterministic_splits(
-        self, tmp_path, mock_admin_db, sample_documents, sample_annotations
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo,
+        sample_documents, sample_annotations
     ):
         """Same seed should produce same splits."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
         results = []
         for _ in range(2):
-            builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-            mock_admin_db.get_documents_by_ids.return_value = sample_documents
-            mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+            builder = DatasetBuilder(
+                datasets_repo=mock_datasets_repo,
+                documents_repo=mock_documents_repo,
+                annotations_repo=mock_annotations_repo,
+                base_dir=tmp_path / "datasets",
+            )
+            mock_documents_repo.get_by_ids.return_value = sample_documents
+            mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
                 sample_annotations.get(str(doc_id), [])
             )
-            mock_admin_db.add_dataset_documents.reset_mock()
-            mock_admin_db.update_dataset_status.reset_mock()
+            mock_datasets_repo.add_documents.reset_mock()
+            mock_datasets_repo.update_status.reset_mock()
 
-            dataset = mock_admin_db.create_dataset.return_value
+            dataset = mock_datasets_repo.create.return_value
             builder.build_dataset(
                 dataset_id=str(dataset.dataset_id),
                 document_ids=[str(d.document_id) for d in sample_documents],
@@ -324,7 +383,7 @@ class TestDatasetBuilder:
                 seed=42,
                 admin_images_dir=tmp_path / "admin_images",
             )
-            call_args = mock_admin_db.add_dataset_documents.call_args
+            call_args = mock_datasets_repo.add_documents.call_args
             docs = call_args[1]["documents"] if "documents" in call_args[1] else call_args[0][1]
             results.append([(d["document_id"], d["split"]) for d in docs])
 
@@ -342,11 +401,18 @@ class TestAssignSplitsByGroup:
         doc.page_count = 1
         return doc
 
-    def test_single_doc_groups_are_distributed(self, tmp_path, mock_admin_db):
+    def test_single_doc_groups_are_distributed(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Documents with unique group_key are distributed across splits."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         # 3 documents, each with unique group_key
         docs = [
@@ -363,11 +429,18 @@ class TestAssignSplitsByGroup:
         assert train_count >= 1
         assert val_count >= 1  # Ensure val is not empty
 
-    def test_null_group_key_treated_as_single_doc_group(self, tmp_path, mock_admin_db):
+    def test_null_group_key_treated_as_single_doc_group(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Documents with null/empty group_key are each treated as independent single-doc groups."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         docs = [
             self._make_mock_doc(uuid4(), group_key=None),
@@ -384,11 +457,18 @@ class TestAssignSplitsByGroup:
         assert train_count >= 1
         assert val_count >= 1
 
-    def test_multi_doc_groups_stay_together(self, tmp_path, mock_admin_db):
+    def test_multi_doc_groups_stay_together(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Documents with same group_key should be assigned to the same split."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         # 6 documents in 2 groups
         docs = [
@@ -410,11 +490,18 @@ class TestAssignSplitsByGroup:
         splits_b = [result[str(d.document_id)] for d in docs[3:]]
         assert len(set(splits_b)) == 1, "All docs in supplier-B should be in same split"
 
-    def test_multi_doc_groups_split_by_ratio(self, tmp_path, mock_admin_db):
+    def test_multi_doc_groups_split_by_ratio(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Multi-doc groups should be split according to train/val/test ratios."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         # 10 groups with 2 docs each
         docs = []
@@ -445,11 +532,18 @@ class TestAssignSplitsByGroup:
         assert split_counts["val"] >= 1
         assert split_counts["val"] <= 3
 
-    def test_mixed_single_and_multi_doc_groups(self, tmp_path, mock_admin_db):
+    def test_mixed_single_and_multi_doc_groups(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Mix of single-doc and multi-doc groups should be handled correctly."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         docs = [
             # Single-doc groups
@@ -476,11 +570,18 @@ class TestAssignSplitsByGroup:
         assert result[str(docs[3].document_id)] == result[str(docs[4].document_id)]
         assert result[str(docs[5].document_id)] == result[str(docs[6].document_id)]
 
-    def test_deterministic_with_seed(self, tmp_path, mock_admin_db):
+    def test_deterministic_with_seed(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Same seed should produce same split assignments."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         docs = [
             self._make_mock_doc(uuid4(), group_key="group-A"),
@@ -496,11 +597,18 @@ class TestAssignSplitsByGroup:
 
         assert result1 == result2
 
-    def test_different_seed_may_produce_different_splits(self, tmp_path, mock_admin_db):
+    def test_different_seed_may_produce_different_splits(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Different seeds should potentially produce different split assignments."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         # Many groups to increase chance of different results
         docs = []
@@ -515,11 +623,18 @@ class TestAssignSplitsByGroup:
         # Results should be different (very likely with 20 groups)
         assert result1 != result2
 
-    def test_all_docs_assigned(self, tmp_path, mock_admin_db):
+    def test_all_docs_assigned(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Every document should be assigned a split."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         docs = [
             self._make_mock_doc(uuid4(), group_key="group-A"),
@@ -535,21 +650,35 @@ class TestAssignSplitsByGroup:
             assert str(doc.document_id) in result
             assert result[str(doc.document_id)] in ["train", "val", "test"]
 
-    def test_empty_documents_list(self, tmp_path, mock_admin_db):
+    def test_empty_documents_list(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Empty document list should return empty result."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         result = builder._assign_splits_by_group([], train_ratio=0.7, val_ratio=0.2, seed=42)
 
         assert result == {}
 
-    def test_only_multi_doc_groups(self, tmp_path, mock_admin_db):
+    def test_only_multi_doc_groups(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """When all groups have multiple docs, splits should follow ratios."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         # 5 groups with 3 docs each
         docs = []
@@ -574,11 +703,18 @@ class TestAssignSplitsByGroup:
         assert split_counts["train"] >= 2
         assert split_counts["train"] <= 4
 
-    def test_only_single_doc_groups(self, tmp_path, mock_admin_db):
+    def test_only_single_doc_groups(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """When all groups have single doc, they are distributed across splits."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
 
         docs = [
             self._make_mock_doc(uuid4(), group_key="unique-1"),
@@ -658,20 +794,26 @@ class TestBuildDatasetWithGroupKey:
         return annotations
 
     def test_build_respects_group_key_splits(
-        self, grouped_documents, grouped_annotations, mock_admin_db
+        self, grouped_documents, grouped_annotations,
+        mock_datasets_repo, mock_documents_repo, mock_annotations_repo
     ):
         """build_dataset should use group_key for split assignment."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
         tmp_path, docs = grouped_documents
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = docs
-        mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = docs
+        mock_annotations_repo.get_for_document.side_effect = lambda doc_id, page_number=None: (
             grouped_annotations.get(str(doc_id), [])
         )
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in docs],
@@ -681,8 +823,8 @@ class TestBuildDatasetWithGroupKey:
             admin_images_dir=tmp_path / "admin_images",
         )
 
-        # Get the document splits from add_dataset_documents call
-        call_args = mock_admin_db.add_dataset_documents.call_args
+        # Get the document splits from add_documents call
+        call_args = mock_datasets_repo.add_documents.call_args
         docs_added = call_args[1]["documents"] if "documents" in call_args[1] else call_args[0][1]
 
         # Build mapping of doc_id -> split
@@ -701,7 +843,9 @@ class TestBuildDatasetWithGroupKey:
         supplier_b_splits = [doc_split_map[doc_id] for doc_id in supplier_b_ids]
         assert len(set(supplier_b_splits)) == 1, "supplier-B docs should be in same split"
 
-    def test_build_with_all_same_group_key(self, tmp_path, mock_admin_db):
+    def test_build_with_all_same_group_key(
+        self, tmp_path, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """All docs with same group_key should go to same split."""
         from inference.web.services.dataset_builder import DatasetBuilder
 
@@ -720,11 +864,16 @@ class TestBuildDatasetWithGroupKey:
             doc.group_key = "same-group"
             docs.append(doc)
 
-        builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
-        mock_admin_db.get_documents_by_ids.return_value = docs
-        mock_admin_db.get_annotations_for_document.return_value = []
+        builder = DatasetBuilder(
+            datasets_repo=mock_datasets_repo,
+            documents_repo=mock_documents_repo,
+            annotations_repo=mock_annotations_repo,
+            base_dir=tmp_path / "datasets",
+        )
+        mock_documents_repo.get_by_ids.return_value = docs
+        mock_annotations_repo.get_for_document.return_value = []
 
-        dataset = mock_admin_db.create_dataset.return_value
+        dataset = mock_datasets_repo.create.return_value
         builder.build_dataset(
             dataset_id=str(dataset.dataset_id),
             document_ids=[str(d.document_id) for d in docs],
@@ -734,7 +883,7 @@ class TestBuildDatasetWithGroupKey:
             admin_images_dir=tmp_path / "admin_images",
         )
 
-        call_args = mock_admin_db.add_dataset_documents.call_args
+        call_args = mock_datasets_repo.add_documents.call_args
         docs_added = call_args[1]["documents"] if "documents" in call_args[1] else call_args[0][1]
 
         splits = [d["split"] for d in docs_added]

tests/web/test_dataset_routes.py

@@ -72,6 +72,36 @@ def _find_endpoint(name: str):
     raise AssertionError(f"Endpoint {name} not found")
 
 
+@pytest.fixture
+def mock_datasets_repo():
+    """Mock DatasetRepository."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_documents_repo():
+    """Mock DocumentRepository."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_annotations_repo():
+    """Mock AnnotationRepository."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_models_repo():
+    """Mock ModelVersionRepository."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_tasks_repo():
+    """Mock TrainingTaskRepository."""
+    return MagicMock()
+
+
 class TestCreateDatasetRoute:
     """Tests for POST /admin/training/datasets."""
 
@@ -80,11 +110,12 @@ class TestCreateDatasetRoute:
         paths = [route.path for route in router.routes]
         assert any("datasets" in p for p in paths)
 
-    def test_create_dataset_calls_builder(self):
+    def test_create_dataset_calls_builder(
+        self, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         fn = _find_endpoint("create_dataset")
 
-        mock_db = MagicMock()
-        mock_db.create_dataset.return_value = _make_dataset(status="building")
+        mock_datasets_repo.create.return_value = _make_dataset(status="building")
 
         mock_builder = MagicMock()
         mock_builder.build_dataset.return_value = {
@@ -101,20 +132,30 @@ class TestCreateDatasetRoute:
         with patch(
             "inference.web.services.dataset_builder.DatasetBuilder",
             return_value=mock_builder,
-        ) as mock_cls:
-            result = asyncio.run(fn(request=request, admin_token=TEST_TOKEN, db=mock_db))
+        ), patch(
+            "inference.web.api.v1.admin.training.datasets.get_storage_helper"
+        ) as mock_storage:
+            mock_storage.return_value.get_datasets_base_path.return_value = "/data/datasets"
+            mock_storage.return_value.get_admin_images_base_path.return_value = "/data/admin_images"
+            result = asyncio.run(fn(
+                request=request,
+                admin_token=TEST_TOKEN,
+                datasets=mock_datasets_repo,
+                docs=mock_documents_repo,
+                annotations=mock_annotations_repo,
+            ))
 
-        mock_db.create_dataset.assert_called_once()
+        mock_datasets_repo.create.assert_called_once()
         mock_builder.build_dataset.assert_called_once()
         assert result.dataset_id == TEST_DATASET_UUID
         assert result.name == "test-dataset"
 
-    def test_create_dataset_fails_with_less_than_10_documents(self):
+    def test_create_dataset_fails_with_less_than_10_documents(
+        self, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Test that creating dataset fails if fewer than 10 documents provided."""
         fn = _find_endpoint("create_dataset")
 
-        mock_db = MagicMock()
-
         # Only 2 documents - should fail
         request = DatasetCreateRequest(
             name="test-dataset",
@@ -124,20 +165,26 @@ class TestCreateDatasetRoute:
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(request=request, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(
+                request=request,
+                admin_token=TEST_TOKEN,
+                datasets=mock_datasets_repo,
+                docs=mock_documents_repo,
+                annotations=mock_annotations_repo,
+            ))
 
         assert exc_info.value.status_code == 400
         assert "Minimum 10 documents required" in exc_info.value.detail
         assert "got 2" in exc_info.value.detail
-        # Ensure DB was never called since validation failed first
-        mock_db.create_dataset.assert_not_called()
+        # Ensure repo was never called since validation failed first
+        mock_datasets_repo.create.assert_not_called()
 
-    def test_create_dataset_fails_with_9_documents(self):
+    def test_create_dataset_fails_with_9_documents(
+        self, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Test boundary condition: 9 documents should fail."""
         fn = _find_endpoint("create_dataset")
 
-        mock_db = MagicMock()
-
         # 9 documents - just under the limit
         request = DatasetCreateRequest(
             name="test-dataset",
@@ -147,17 +194,24 @@ class TestCreateDatasetRoute:
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(request=request, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(
+                request=request,
+                admin_token=TEST_TOKEN,
+                datasets=mock_datasets_repo,
+                docs=mock_documents_repo,
+                annotations=mock_annotations_repo,
+            ))
 
         assert exc_info.value.status_code == 400
         assert "Minimum 10 documents required" in exc_info.value.detail
 
-    def test_create_dataset_succeeds_with_exactly_10_documents(self):
+    def test_create_dataset_succeeds_with_exactly_10_documents(
+        self, mock_datasets_repo, mock_documents_repo, mock_annotations_repo
+    ):
         """Test boundary condition: exactly 10 documents should succeed."""
         fn = _find_endpoint("create_dataset")
 
-        mock_db = MagicMock()
-        mock_db.create_dataset.return_value = _make_dataset(status="building")
+        mock_datasets_repo.create.return_value = _make_dataset(status="building")
 
         mock_builder = MagicMock()
 
@@ -170,25 +224,40 @@ class TestCreateDatasetRoute:
         with patch(
             "inference.web.services.dataset_builder.DatasetBuilder",
             return_value=mock_builder,
-        ):
-            result = asyncio.run(fn(request=request, admin_token=TEST_TOKEN, db=mock_db))
+        ), patch(
+            "inference.web.api.v1.admin.training.datasets.get_storage_helper"
+        ) as mock_storage:
+            mock_storage.return_value.get_datasets_base_path.return_value = "/data/datasets"
+            mock_storage.return_value.get_admin_images_base_path.return_value = "/data/admin_images"
+            result = asyncio.run(fn(
+                request=request,
+                admin_token=TEST_TOKEN,
+                datasets=mock_datasets_repo,
+                docs=mock_documents_repo,
+                annotations=mock_annotations_repo,
+            ))
 
-        mock_db.create_dataset.assert_called_once()
+        mock_datasets_repo.create.assert_called_once()
         assert result.dataset_id == TEST_DATASET_UUID
 
 
 class TestListDatasetsRoute:
     """Tests for GET /admin/training/datasets."""
 
-    def test_list_datasets(self):
+    def test_list_datasets(self, mock_datasets_repo):
         fn = _find_endpoint("list_datasets")
 
-        mock_db = MagicMock()
-        mock_db.get_datasets.return_value = ([_make_dataset()], 1)
+        mock_datasets_repo.get_paginated.return_value = ([_make_dataset()], 1)
         # Mock the active training tasks lookup to return empty dict
-        mock_db.get_active_training_tasks_for_datasets.return_value = {}
+        mock_datasets_repo.get_active_training_tasks.return_value = {}
 
-        result = asyncio.run(fn(admin_token=TEST_TOKEN, db=mock_db, status=None, limit=20, offset=0))
+        result = asyncio.run(fn(
+            admin_token=TEST_TOKEN,
+            datasets_repo=mock_datasets_repo,
+            status=None,
+            limit=20,
+            offset=0,
+        ))
 
         assert result.total == 1
         assert len(result.datasets) == 1
@@ -198,82 +267,103 @@ class TestListDatasetsRoute:
 class TestGetDatasetRoute:
     """Tests for GET /admin/training/datasets/{dataset_id}."""
 
-    def test_get_dataset_returns_detail(self):
+    def test_get_dataset_returns_detail(self, mock_datasets_repo):
         fn = _find_endpoint("get_dataset")
 
-        mock_db = MagicMock()
-        mock_db.get_dataset.return_value = _make_dataset()
-        mock_db.get_dataset_documents.return_value = [
+        mock_datasets_repo.get.return_value = _make_dataset()
+        mock_datasets_repo.get_documents.return_value = [
             _make_dataset_doc(TEST_DOC_UUID_1, "train"),
             _make_dataset_doc(TEST_DOC_UUID_2, "val"),
         ]
 
-        result = asyncio.run(fn(dataset_id=TEST_DATASET_UUID, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(
+            dataset_id=TEST_DATASET_UUID,
+            admin_token=TEST_TOKEN,
+            datasets_repo=mock_datasets_repo,
+        ))
 
         assert result.dataset_id == TEST_DATASET_UUID
         assert len(result.documents) == 2
 
-    def test_get_dataset_not_found(self):
+    def test_get_dataset_not_found(self, mock_datasets_repo):
         fn = _find_endpoint("get_dataset")
 
-        mock_db = MagicMock()
-        mock_db.get_dataset.return_value = None
+        mock_datasets_repo.get.return_value = None
 
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(dataset_id=TEST_DATASET_UUID, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(
+                dataset_id=TEST_DATASET_UUID,
+                admin_token=TEST_TOKEN,
+                datasets_repo=mock_datasets_repo,
+            ))
         assert exc_info.value.status_code == 404
 
 
 class TestDeleteDatasetRoute:
     """Tests for DELETE /admin/training/datasets/{dataset_id}."""
 
-    def test_delete_dataset(self):
+    def test_delete_dataset(self, mock_datasets_repo):
         fn = _find_endpoint("delete_dataset")
 
-        mock_db = MagicMock()
-        mock_db.get_dataset.return_value = _make_dataset(dataset_path=None)
+        mock_datasets_repo.get.return_value = _make_dataset(dataset_path=None)
 
-        result = asyncio.run(fn(dataset_id=TEST_DATASET_UUID, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(
+            dataset_id=TEST_DATASET_UUID,
+            admin_token=TEST_TOKEN,
+            datasets_repo=mock_datasets_repo,
+        ))
 
-        mock_db.delete_dataset.assert_called_once_with(TEST_DATASET_UUID)
+        mock_datasets_repo.delete.assert_called_once_with(TEST_DATASET_UUID)
         assert result["message"] == "Dataset deleted"
 
 
 class TestTrainFromDatasetRoute:
     """Tests for POST /admin/training/datasets/{dataset_id}/train."""
 
-    def test_train_from_ready_dataset(self):
+    def test_train_from_ready_dataset(self, mock_datasets_repo, mock_models_repo, mock_tasks_repo):
         fn = _find_endpoint("train_from_dataset")
 
-        mock_db = MagicMock()
-        mock_db.get_dataset.return_value = _make_dataset(status="ready")
-        mock_db.create_training_task.return_value = TEST_TASK_UUID
+        mock_datasets_repo.get.return_value = _make_dataset(status="ready")
+        mock_tasks_repo.create.return_value = TEST_TASK_UUID
 
         request = DatasetTrainRequest(name="train-from-dataset", config=TrainingConfig())
 
-        result = asyncio.run(fn(dataset_id=TEST_DATASET_UUID, request=request, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(
+            dataset_id=TEST_DATASET_UUID,
+            request=request,
+            admin_token=TEST_TOKEN,
+            datasets_repo=mock_datasets_repo,
+            models=mock_models_repo,
+            tasks=mock_tasks_repo,
+        ))
 
         assert result.task_id == TEST_TASK_UUID
         assert result.status == TrainingStatus.PENDING
-        mock_db.create_training_task.assert_called_once()
+        mock_tasks_repo.create.assert_called_once()
 
-    def test_train_from_building_dataset_fails(self):
+    def test_train_from_building_dataset_fails(self, mock_datasets_repo, mock_models_repo, mock_tasks_repo):
         fn = _find_endpoint("train_from_dataset")
 
-        mock_db = MagicMock()
-        mock_db.get_dataset.return_value = _make_dataset(status="building")
+        mock_datasets_repo.get.return_value = _make_dataset(status="building")
 
         request = DatasetTrainRequest(name="train-from-dataset", config=TrainingConfig())
 
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(dataset_id=TEST_DATASET_UUID, request=request, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(
+                dataset_id=TEST_DATASET_UUID,
+                request=request,
+                admin_token=TEST_TOKEN,
+                datasets_repo=mock_datasets_repo,
+                models=mock_models_repo,
+                tasks=mock_tasks_repo,
+            ))
         assert exc_info.value.status_code == 400
 
-    def test_incremental_training_with_base_model(self):
+    def test_incremental_training_with_base_model(self, mock_datasets_repo, mock_models_repo, mock_tasks_repo):
         """Test training with base_model_version_id for incremental training."""
         fn = _find_endpoint("train_from_dataset")
 
@@ -281,22 +371,28 @@ class TestTrainFromDatasetRoute:
         mock_model_version.model_path = "runs/train/invoice_fields/weights/best.pt"
         mock_model_version.version = "1.0.0"
 
-        mock_db = MagicMock()
-        mock_db.get_dataset.return_value = _make_dataset(status="ready")
-        mock_db.get_model_version.return_value = mock_model_version
-        mock_db.create_training_task.return_value = TEST_TASK_UUID
+        mock_datasets_repo.get.return_value = _make_dataset(status="ready")
+        mock_models_repo.get.return_value = mock_model_version
+        mock_tasks_repo.create.return_value = TEST_TASK_UUID
 
         base_model_uuid = "550e8400-e29b-41d4-a716-446655440099"
         config = TrainingConfig(base_model_version_id=base_model_uuid)
         request = DatasetTrainRequest(name="incremental-train", config=config)
 
-        result = asyncio.run(fn(dataset_id=TEST_DATASET_UUID, request=request, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(
+            dataset_id=TEST_DATASET_UUID,
+            request=request,
+            admin_token=TEST_TOKEN,
+            datasets_repo=mock_datasets_repo,
+            models=mock_models_repo,
+            tasks=mock_tasks_repo,
+        ))
 
         # Verify model version was looked up
-        mock_db.get_model_version.assert_called_once_with(base_model_uuid)
+        mock_models_repo.get.assert_called_once_with(base_model_uuid)
 
         # Verify task was created with finetune type
-        call_kwargs = mock_db.create_training_task.call_args[1]
+        call_kwargs = mock_tasks_repo.create.call_args[1]
         assert call_kwargs["task_type"] == "finetune"
         assert call_kwargs["config"]["base_model_path"] == "runs/train/invoice_fields/weights/best.pt"
         assert call_kwargs["config"]["base_model_version"] == "1.0.0"
@@ -304,13 +400,14 @@ class TestTrainFromDatasetRoute:
         assert result.task_id == TEST_TASK_UUID
         assert "Incremental training" in result.message
 
-    def test_incremental_training_with_invalid_base_model_fails(self):
+    def test_incremental_training_with_invalid_base_model_fails(
+        self, mock_datasets_repo, mock_models_repo, mock_tasks_repo
+    ):
         """Test that training fails if base_model_version_id doesn't exist."""
         fn = _find_endpoint("train_from_dataset")
 
-        mock_db = MagicMock()
-        mock_db.get_dataset.return_value = _make_dataset(status="ready")
-        mock_db.get_model_version.return_value = None
+        mock_datasets_repo.get.return_value = _make_dataset(status="ready")
+        mock_models_repo.get.return_value = None
 
         base_model_uuid = "550e8400-e29b-41d4-a716-446655440099"
         config = TrainingConfig(base_model_version_id=base_model_uuid)
@@ -319,6 +416,13 @@ class TestTrainFromDatasetRoute:
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(dataset_id=TEST_DATASET_UUID, request=request, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(
+                dataset_id=TEST_DATASET_UUID,
+                request=request,
+                admin_token=TEST_TOKEN,
+                datasets_repo=mock_datasets_repo,
+                models=mock_models_repo,
+                tasks=mock_tasks_repo,
+            ))
         assert exc_info.value.status_code == 404
         assert "Base model version not found" in exc_info.value.detail

tests/web/test_dataset_training_status.py

@@ -3,7 +3,7 @@ Tests for dataset training status feature.
 
 Tests cover:
 1. Database model fields (training_status, active_training_task_id)
-2. AdminDB update_dataset_training_status method
+2. DatasetRepository update_training_status method
 3. API response includes training status fields
 4. Scheduler updates dataset status during training lifecycle
 """
@@ -56,12 +56,12 @@ class TestTrainingDatasetModel:
 
 
 # =============================================================================
-# Test AdminDB Methods
+# Test DatasetRepository Methods
 # =============================================================================
 
 
-class TestAdminDBDatasetTrainingStatus:
-    """Tests for AdminDB.update_dataset_training_status method."""
+class TestDatasetRepositoryTrainingStatus:
+    """Tests for DatasetRepository.update_training_status method."""
 
     @pytest.fixture
     def mock_session(self):
@@ -69,8 +69,8 @@ class TestAdminDBDatasetTrainingStatus:
         session = MagicMock()
         return session
 
-    def test_update_dataset_training_status_sets_status(self, mock_session):
-        """update_dataset_training_status should set training_status."""
+    def test_update_training_status_sets_status(self, mock_session):
+        """update_training_status should set training_status."""
         from inference.data.admin_models import TrainingDataset
 
         dataset_id = uuid4()
@@ -81,13 +81,13 @@ class TestAdminDBDatasetTrainingStatus:
         )
         mock_session.get.return_value = dataset
 
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
             mock_ctx.return_value.__enter__.return_value = mock_session
 
-            from inference.data.admin_db import AdminDB
+            from inference.data.repositories import DatasetRepository
 
-            db = AdminDB()
-            db.update_dataset_training_status(
+            repo = DatasetRepository()
+            repo.update_training_status(
                 dataset_id=str(dataset_id),
                 training_status="running",
             )
@@ -96,8 +96,8 @@ class TestAdminDBDatasetTrainingStatus:
         mock_session.add.assert_called_once_with(dataset)
         mock_session.commit.assert_called_once()
 
-    def test_update_dataset_training_status_sets_task_id(self, mock_session):
-        """update_dataset_training_status should set active_training_task_id."""
+    def test_update_training_status_sets_task_id(self, mock_session):
+        """update_training_status should set active_training_task_id."""
         from inference.data.admin_models import TrainingDataset
 
         dataset_id = uuid4()
@@ -109,13 +109,13 @@ class TestAdminDBDatasetTrainingStatus:
         )
         mock_session.get.return_value = dataset
 
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
             mock_ctx.return_value.__enter__.return_value = mock_session
 
-            from inference.data.admin_db import AdminDB
+            from inference.data.repositories import DatasetRepository
 
-            db = AdminDB()
-            db.update_dataset_training_status(
+            repo = DatasetRepository()
+            repo.update_training_status(
                 dataset_id=str(dataset_id),
                 training_status="running",
                 active_training_task_id=str(task_id),
@@ -123,10 +123,10 @@ class TestAdminDBDatasetTrainingStatus:
 
         assert dataset.active_training_task_id == task_id
 
-    def test_update_dataset_training_status_updates_main_status_on_complete(
+    def test_update_training_status_updates_main_status_on_complete(
         self, mock_session
     ):
-        """update_dataset_training_status should update main status to 'trained' when completed."""
+        """update_training_status should update main status to 'trained' when completed."""
         from inference.data.admin_models import TrainingDataset
 
         dataset_id = uuid4()
@@ -137,13 +137,13 @@ class TestAdminDBDatasetTrainingStatus:
         )
         mock_session.get.return_value = dataset
 
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
             mock_ctx.return_value.__enter__.return_value = mock_session
 
-            from inference.data.admin_db import AdminDB
+            from inference.data.repositories import DatasetRepository
 
-            db = AdminDB()
-            db.update_dataset_training_status(
+            repo = DatasetRepository()
+            repo.update_training_status(
                 dataset_id=str(dataset_id),
                 training_status="completed",
                 update_main_status=True,
@@ -152,10 +152,10 @@ class TestAdminDBDatasetTrainingStatus:
         assert dataset.status == "trained"
         assert dataset.training_status == "completed"
 
-    def test_update_dataset_training_status_clears_task_id_on_complete(
+    def test_update_training_status_clears_task_id_on_complete(
         self, mock_session
     ):
-        """update_dataset_training_status should clear task_id when training completes."""
+        """update_training_status should clear task_id when training completes."""
         from inference.data.admin_models import TrainingDataset
 
         dataset_id = uuid4()
@@ -169,13 +169,13 @@ class TestAdminDBDatasetTrainingStatus:
         )
         mock_session.get.return_value = dataset
 
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
             mock_ctx.return_value.__enter__.return_value = mock_session
 
-            from inference.data.admin_db import AdminDB
+            from inference.data.repositories import DatasetRepository
 
-            db = AdminDB()
-            db.update_dataset_training_status(
+            repo = DatasetRepository()
+            repo.update_training_status(
                 dataset_id=str(dataset_id),
                 training_status="completed",
                 active_training_task_id=None,
@@ -183,18 +183,18 @@ class TestAdminDBDatasetTrainingStatus:
 
         assert dataset.active_training_task_id is None
 
-    def test_update_dataset_training_status_handles_missing_dataset(self, mock_session):
-        """update_dataset_training_status should handle missing dataset gracefully."""
+    def test_update_training_status_handles_missing_dataset(self, mock_session):
+        """update_training_status should handle missing dataset gracefully."""
         mock_session.get.return_value = None
 
-        with patch("inference.data.admin_db.get_session_context") as mock_ctx:
+        with patch("inference.data.repositories.dataset_repository.get_session_context") as mock_ctx:
             mock_ctx.return_value.__enter__.return_value = mock_session
 
-            from inference.data.admin_db import AdminDB
+            from inference.data.repositories import DatasetRepository
 
-            db = AdminDB()
+            repo = DatasetRepository()
             # Should not raise
-            db.update_dataset_training_status(
+            repo.update_training_status(
                 dataset_id=str(uuid4()),
                 training_status="running",
             )
@@ -275,19 +275,24 @@ class TestSchedulerDatasetStatusUpdates:
     """Tests for scheduler updating dataset status during training."""
 
     @pytest.fixture
-    def mock_db(self):
-        """Create mock AdminDB."""
+    def mock_datasets_repo(self):
+        """Create mock DatasetRepository."""
         mock = MagicMock()
-        mock.get_dataset.return_value = MagicMock(
+        mock.get.return_value = MagicMock(
             dataset_id=uuid4(),
             name="test-dataset",
             dataset_path="/path/to/dataset",
             total_images=100,
         )
-        mock.get_pending_training_tasks.return_value = []
         return mock
 
-    def test_scheduler_sets_running_status_on_task_start(self, mock_db):
+    @pytest.fixture
+    def mock_training_tasks_repo(self):
+        """Create mock TrainingTaskRepository."""
+        mock = MagicMock()
+        return mock
+
+    def test_scheduler_sets_running_status_on_task_start(self, mock_datasets_repo, mock_training_tasks_repo):
         """Scheduler should set dataset training_status to 'running' when task starts."""
         from inference.web.core.scheduler import TrainingScheduler
 
@@ -295,7 +300,8 @@ class TestSchedulerDatasetStatusUpdates:
             mock_train.return_value = {"model_path": "/path/to/model.pt", "metrics": {}}
 
             scheduler = TrainingScheduler()
-            scheduler._db = mock_db
+            scheduler._datasets = mock_datasets_repo
+            scheduler._training_tasks = mock_training_tasks_repo
 
             task_id = str(uuid4())
             dataset_id = str(uuid4())
@@ -311,8 +317,8 @@ class TestSchedulerDatasetStatusUpdates:
                 pass  # Expected to fail in test environment
 
             # Check that training status was updated to running
-            mock_db.update_dataset_training_status.assert_called()
-            first_call = mock_db.update_dataset_training_status.call_args_list[0]
+            mock_datasets_repo.update_training_status.assert_called()
+            first_call = mock_datasets_repo.update_training_status.call_args_list[0]
             assert first_call.kwargs["training_status"] == "running"
             assert first_call.kwargs["active_training_task_id"] == task_id
 

tests/web/test_document_category_api.py

@@ -45,10 +45,10 @@ class TestDocumentListFilterByCategory:
     """Tests for filtering documents by category."""
 
     @pytest.fixture
-    def mock_admin_db(self):
-        """Create mock AdminDB."""
-        db = MagicMock()
-        db.is_valid_admin_token.return_value = True
+    def mock_document_repo(self):
+        """Create mock DocumentRepository."""
+        repo = MagicMock()
+        repo.is_valid.return_value = True
 
         # Mock documents with different categories
         invoice_doc = MagicMock()
@@ -61,11 +61,11 @@ class TestDocumentListFilterByCategory:
         letter_doc.category = "letter"
         letter_doc.filename = "letter1.pdf"
 
-        db.get_documents.return_value = ([invoice_doc], 1)
-        db.get_document_categories.return_value = ["invoice", "letter", "receipt"]
-        return db
+        repo.get_paginated.return_value = ([invoice_doc], 1)
+        repo.get_categories.return_value = ["invoice", "letter", "receipt"]
+        return repo
 
-    def test_list_documents_accepts_category_filter(self, mock_admin_db):
+    def test_list_documents_accepts_category_filter(self, mock_document_repo):
         """Test list documents endpoint accepts category query parameter."""
         # The endpoint should accept ?category=invoice parameter
         # This test verifies the schema/query parameter exists
@@ -74,9 +74,9 @@ class TestDocumentListFilterByCategory:
         # Schema should work with category filter applied
         assert DocumentListResponse is not None
 
-    def test_get_document_categories_from_db(self, mock_admin_db):
-        """Test fetching unique categories from database."""
-        categories = mock_admin_db.get_document_categories()
+    def test_get_document_categories_from_repo(self, mock_document_repo):
+        """Test fetching unique categories from repository."""
+        categories = mock_document_repo.get_categories()
         assert "invoice" in categories
         assert "letter" in categories
         assert len(categories) == 3
@@ -122,24 +122,24 @@ class TestDocumentUploadWithCategory:
         assert response.category == "invoice"
 
 
-class TestAdminDBCategoryMethods:
-    """Tests for AdminDB category-related methods."""
+class TestDocumentRepositoryCategoryMethods:
+    """Tests for DocumentRepository category-related methods."""
 
-    def test_get_document_categories_method_exists(self):
-        """Test AdminDB has get_document_categories method."""
-        from inference.data.admin_db import AdminDB
+    def test_get_categories_method_exists(self):
+        """Test DocumentRepository has get_categories method."""
+        from inference.data.repositories import DocumentRepository
 
-        db = AdminDB()
-        assert hasattr(db, "get_document_categories")
+        repo = DocumentRepository()
+        assert hasattr(repo, "get_categories")
 
-    def test_get_documents_accepts_category_filter(self):
-        """Test get_documents_by_token method accepts category parameter."""
-        from inference.data.admin_db import AdminDB
+    def test_get_paginated_accepts_category_filter(self):
+        """Test get_paginated method accepts category parameter."""
+        from inference.data.repositories import DocumentRepository
         import inspect
 
-        db = AdminDB()
+        repo = DocumentRepository()
         # Check the method exists and accepts category parameter
-        method = getattr(db, "get_documents_by_token", None)
+        method = getattr(repo, "get_paginated", None)
         assert callable(method)
 
         # Check category is in the method signature
@@ -150,12 +150,12 @@ class TestAdminDBCategoryMethods:
 class TestUpdateDocumentCategory:
     """Tests for updating document category."""
 
-    def test_update_document_category_method_exists(self):
-        """Test AdminDB has method to update document category."""
-        from inference.data.admin_db import AdminDB
+    def test_update_category_method_exists(self):
+        """Test DocumentRepository has method to update document category."""
+        from inference.data.repositories import DocumentRepository
 
-        db = AdminDB()
-        assert hasattr(db, "update_document_category")
+        repo = DocumentRepository()
+        assert hasattr(repo, "update_category")
 
     def test_update_request_schema(self):
         """Test DocumentUpdateRequest can update category."""

tests/web/test_model_versions.py

@@ -63,6 +63,12 @@ def _find_endpoint(name: str):
     raise AssertionError(f"Endpoint {name} not found")
 
 
+@pytest.fixture
+def mock_models_repo():
+    """Mock ModelVersionRepository."""
+    return MagicMock()
+
+
 class TestModelVersionRouterRegistration:
     """Tests that model version endpoints are registered."""
 
@@ -91,11 +97,10 @@ class TestModelVersionRouterRegistration:
 class TestCreateModelVersionRoute:
     """Tests for POST /admin/training/models."""
 
-    def test_create_model_version(self):
+    def test_create_model_version(self, mock_models_repo):
         fn = _find_endpoint("create_model_version")
 
-        mock_db = MagicMock()
-        mock_db.create_model_version.return_value = _make_model_version()
+        mock_models_repo.create.return_value = _make_model_version()
 
         request = ModelVersionCreateRequest(
             version="1.0.0",
@@ -106,18 +111,17 @@ class TestCreateModelVersionRoute:
             document_count=100,
         )
 
-        result = asyncio.run(fn(request=request, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(request=request, admin_token=TEST_TOKEN, models=mock_models_repo))
 
-        mock_db.create_model_version.assert_called_once()
+        mock_models_repo.create.assert_called_once()
         assert result.version_id == TEST_VERSION_UUID
         assert result.status == "inactive"
         assert result.message == "Model version created successfully"
 
-    def test_create_model_version_with_task_and_dataset(self):
+    def test_create_model_version_with_task_and_dataset(self, mock_models_repo):
         fn = _find_endpoint("create_model_version")
 
-        mock_db = MagicMock()
-        mock_db.create_model_version.return_value = _make_model_version()
+        mock_models_repo.create.return_value = _make_model_version()
 
         request = ModelVersionCreateRequest(
             version="1.0.0",
@@ -127,9 +131,9 @@ class TestCreateModelVersionRoute:
             dataset_id=TEST_DATASET_UUID,
         )
 
-        result = asyncio.run(fn(request=request, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(request=request, admin_token=TEST_TOKEN, models=mock_models_repo))
 
-        call_kwargs = mock_db.create_model_version.call_args[1]
+        call_kwargs = mock_models_repo.create.call_args[1]
         assert call_kwargs["task_id"] == TEST_TASK_UUID
         assert call_kwargs["dataset_id"] == TEST_DATASET_UUID
 
@@ -137,30 +141,28 @@ class TestCreateModelVersionRoute:
 class TestListModelVersionsRoute:
     """Tests for GET /admin/training/models."""
 
-    def test_list_model_versions(self):
+    def test_list_model_versions(self, mock_models_repo):
         fn = _find_endpoint("list_model_versions")
 
-        mock_db = MagicMock()
-        mock_db.get_model_versions.return_value = (
+        mock_models_repo.get_paginated.return_value = (
             [_make_model_version(), _make_model_version(version_id=UUID(TEST_VERSION_UUID_2), version="1.1.0")],
             2,
         )
 
-        result = asyncio.run(fn(admin_token=TEST_TOKEN, db=mock_db, status=None, limit=20, offset=0))
+        result = asyncio.run(fn(admin_token=TEST_TOKEN, models=mock_models_repo, status=None, limit=20, offset=0))
 
         assert result.total == 2
         assert len(result.models) == 2
         assert result.models[0].version == "1.0.0"
 
-    def test_list_model_versions_with_status_filter(self):
+    def test_list_model_versions_with_status_filter(self, mock_models_repo):
         fn = _find_endpoint("list_model_versions")
 
-        mock_db = MagicMock()
-        mock_db.get_model_versions.return_value = ([_make_model_version(status="active", is_active=True)], 1)
+        mock_models_repo.get_paginated.return_value = ([_make_model_version(status="active", is_active=True)], 1)
 
-        result = asyncio.run(fn(admin_token=TEST_TOKEN, db=mock_db, status="active", limit=20, offset=0))
+        result = asyncio.run(fn(admin_token=TEST_TOKEN, models=mock_models_repo, status="active", limit=20, offset=0))
 
-        mock_db.get_model_versions.assert_called_once_with(status="active", limit=20, offset=0)
+        mock_models_repo.get_paginated.assert_called_once_with(status="active", limit=20, offset=0)
         assert result.total == 1
         assert result.models[0].status == "active"
 
@@ -168,25 +170,23 @@ class TestListModelVersionsRoute:
 class TestGetActiveModelRoute:
     """Tests for GET /admin/training/models/active."""
 
-    def test_get_active_model_when_exists(self):
+    def test_get_active_model_when_exists(self, mock_models_repo):
         fn = _find_endpoint("get_active_model")
 
-        mock_db = MagicMock()
-        mock_db.get_active_model_version.return_value = _make_model_version(status="active", is_active=True)
+        mock_models_repo.get_active.return_value = _make_model_version(status="active", is_active=True)
 
-        result = asyncio.run(fn(admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(admin_token=TEST_TOKEN, models=mock_models_repo))
 
         assert result.has_active_model is True
         assert result.model is not None
         assert result.model.is_active is True
 
-    def test_get_active_model_when_none(self):
+    def test_get_active_model_when_none(self, mock_models_repo):
         fn = _find_endpoint("get_active_model")
 
-        mock_db = MagicMock()
-        mock_db.get_active_model_version.return_value = None
+        mock_models_repo.get_active.return_value = None
 
-        result = asyncio.run(fn(admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(admin_token=TEST_TOKEN, models=mock_models_repo))
 
         assert result.has_active_model is False
         assert result.model is None
@@ -195,46 +195,43 @@ class TestGetActiveModelRoute:
 class TestGetModelVersionRoute:
     """Tests for GET /admin/training/models/{version_id}."""
 
-    def test_get_model_version(self):
+    def test_get_model_version(self, mock_models_repo):
         fn = _find_endpoint("get_model_version")
 
-        mock_db = MagicMock()
-        mock_db.get_model_version.return_value = _make_model_version()
+        mock_models_repo.get.return_value = _make_model_version()
 
-        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
 
         assert result.version_id == TEST_VERSION_UUID
         assert result.version == "1.0.0"
         assert result.name == "test-model-v1"
         assert result.metrics_mAP == 0.935
 
-    def test_get_model_version_not_found(self):
+    def test_get_model_version_not_found(self, mock_models_repo):
         fn = _find_endpoint("get_model_version")
 
-        mock_db = MagicMock()
-        mock_db.get_model_version.return_value = None
+        mock_models_repo.get.return_value = None
 
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
         assert exc_info.value.status_code == 404
 
 
 class TestUpdateModelVersionRoute:
     """Tests for PATCH /admin/training/models/{version_id}."""
 
-    def test_update_model_version(self):
+    def test_update_model_version(self, mock_models_repo):
         fn = _find_endpoint("update_model_version")
 
-        mock_db = MagicMock()
-        mock_db.update_model_version.return_value = _make_model_version(name="updated-name")
+        mock_models_repo.update.return_value = _make_model_version(name="updated-name")
 
         request = ModelVersionUpdateRequest(name="updated-name", description="Updated description")
 
-        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, request=request, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, request=request, admin_token=TEST_TOKEN, models=mock_models_repo))
 
-        mock_db.update_model_version.assert_called_once_with(
+        mock_models_repo.update.assert_called_once_with(
             version_id=TEST_VERSION_UUID,
             name="updated-name",
             description="Updated description",
@@ -242,45 +239,42 @@ class TestUpdateModelVersionRoute:
         )
         assert result.message == "Model version updated successfully"
 
-    def test_update_model_version_not_found(self):
+    def test_update_model_version_not_found(self, mock_models_repo):
         fn = _find_endpoint("update_model_version")
 
-        mock_db = MagicMock()
-        mock_db.update_model_version.return_value = None
+        mock_models_repo.update.return_value = None
 
         request = ModelVersionUpdateRequest(name="updated-name")
 
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(version_id=TEST_VERSION_UUID, request=request, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(version_id=TEST_VERSION_UUID, request=request, admin_token=TEST_TOKEN, models=mock_models_repo))
         assert exc_info.value.status_code == 404
 
 
 class TestActivateModelVersionRoute:
     """Tests for POST /admin/training/models/{version_id}/activate."""
 
-    def test_activate_model_version(self):
+    def test_activate_model_version(self, mock_models_repo):
         fn = _find_endpoint("activate_model_version")
 
-        mock_db = MagicMock()
-        mock_db.activate_model_version.return_value = _make_model_version(status="active", is_active=True)
+        mock_models_repo.activate.return_value = _make_model_version(status="active", is_active=True)
 
         # Create mock request with app state
         mock_request = MagicMock()
         mock_request.app.state.inference_service = None
 
-        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, request=mock_request, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, request=mock_request, admin_token=TEST_TOKEN, models=mock_models_repo))
 
-        mock_db.activate_model_version.assert_called_once_with(TEST_VERSION_UUID)
+        mock_models_repo.activate.assert_called_once_with(TEST_VERSION_UUID)
         assert result.status == "active"
         assert result.message == "Model version activated for inference"
 
-    def test_activate_model_version_not_found(self):
+    def test_activate_model_version_not_found(self, mock_models_repo):
         fn = _find_endpoint("activate_model_version")
 
-        mock_db = MagicMock()
-        mock_db.activate_model_version.return_value = None
+        mock_models_repo.activate.return_value = None
 
         # Create mock request with app state
         mock_request = MagicMock()
@@ -289,88 +283,82 @@ class TestActivateModelVersionRoute:
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(version_id=TEST_VERSION_UUID, request=mock_request, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(version_id=TEST_VERSION_UUID, request=mock_request, admin_token=TEST_TOKEN, models=mock_models_repo))
         assert exc_info.value.status_code == 404
 
 
 class TestDeactivateModelVersionRoute:
     """Tests for POST /admin/training/models/{version_id}/deactivate."""
 
-    def test_deactivate_model_version(self):
+    def test_deactivate_model_version(self, mock_models_repo):
         fn = _find_endpoint("deactivate_model_version")
 
-        mock_db = MagicMock()
-        mock_db.deactivate_model_version.return_value = _make_model_version(status="inactive", is_active=False)
+        mock_models_repo.deactivate.return_value = _make_model_version(status="inactive", is_active=False)
 
-        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
 
         assert result.status == "inactive"
         assert result.message == "Model version deactivated"
 
-    def test_deactivate_model_version_not_found(self):
+    def test_deactivate_model_version_not_found(self, mock_models_repo):
         fn = _find_endpoint("deactivate_model_version")
 
-        mock_db = MagicMock()
-        mock_db.deactivate_model_version.return_value = None
+        mock_models_repo.deactivate.return_value = None
 
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
         assert exc_info.value.status_code == 404
 
 
 class TestArchiveModelVersionRoute:
     """Tests for POST /admin/training/models/{version_id}/archive."""
 
-    def test_archive_model_version(self):
+    def test_archive_model_version(self, mock_models_repo):
         fn = _find_endpoint("archive_model_version")
 
-        mock_db = MagicMock()
-        mock_db.archive_model_version.return_value = _make_model_version(status="archived")
+        mock_models_repo.archive.return_value = _make_model_version(status="archived")
 
-        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
 
         assert result.status == "archived"
         assert result.message == "Model version archived"
 
-    def test_archive_active_model_fails(self):
+    def test_archive_active_model_fails(self, mock_models_repo):
         fn = _find_endpoint("archive_model_version")
 
-        mock_db = MagicMock()
-        mock_db.archive_model_version.return_value = None
+        mock_models_repo.archive.return_value = None
 
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
         assert exc_info.value.status_code == 400
 
 
 class TestDeleteModelVersionRoute:
     """Tests for DELETE /admin/training/models/{version_id}."""
 
-    def test_delete_model_version(self):
+    def test_delete_model_version(self, mock_models_repo):
         fn = _find_endpoint("delete_model_version")
 
-        mock_db = MagicMock()
-        mock_db.delete_model_version.return_value = True
+        mock_models_repo.delete.return_value = True
 
-        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+        result = asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
 
-        mock_db.delete_model_version.assert_called_once_with(TEST_VERSION_UUID)
+        mock_models_repo.delete.assert_called_once_with(TEST_VERSION_UUID)
         assert result["message"] == "Model version deleted"
 
-    def test_delete_active_model_fails(self):
+    def test_delete_active_model_fails(self, mock_models_repo):
         fn = _find_endpoint("delete_model_version")
 
-        mock_db = MagicMock()
-        mock_db.delete_model_version.return_value = False
+        mock_models_repo.delete.return_value = False
 
         from fastapi import HTTPException
 
         with pytest.raises(HTTPException) as exc_info:
-            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, db=mock_db))
+            asyncio.run(fn(version_id=TEST_VERSION_UUID, admin_token=TEST_TOKEN, models=mock_models_repo))
         assert exc_info.value.status_code == 400
 
 

tests/web/test_training_phase4.py

@@ -10,7 +10,13 @@ from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
 from inference.web.api.v1.admin.training import create_training_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,
+    get_annotation_repository,
+    get_training_task_repository,
+    get_model_version_repository,
+)
 
 
 class MockTrainingTask:
@@ -128,19 +134,17 @@ class MockModelVersion:
         self.updated_at = kwargs.get('updated_at', datetime.utcnow())
 
 
-class MockAdminDB:
-    """Mock AdminDB for testing Phase 4."""
+class MockDocumentRepository:
+    """Mock DocumentRepository for testing Phase 4."""
 
     def __init__(self):
         self.documents = {}
-        self.annotations = {}
-        self.training_tasks = {}
-        self.training_links = {}
-        self.model_versions = {}
+        self.annotations = {}  # Shared reference for filtering
+        self.training_links = {}  # Shared reference for filtering
 
-    def get_documents_for_training(
+    def get_for_training(
         self,
-        admin_token,
+        admin_token=None,
         status="labeled",
         has_annotations=True,
         min_annotation_count=None,
@@ -173,17 +177,28 @@ class MockAdminDB:
         total = len(filtered)
         return filtered[offset:offset+limit], total
 
-    def get_annotations_for_document(self, document_id):
+
+class MockAnnotationRepository:
+    """Mock AnnotationRepository for testing Phase 4."""
+
+    def __init__(self):
+        self.annotations = {}
+
+    def get_for_document(self, document_id, page_number=None):
         """Get annotations for document."""
         return self.annotations.get(str(document_id), [])
 
-    def get_document_training_tasks(self, document_id):
-        """Get training tasks that used this document."""
-        return self.training_links.get(str(document_id), [])
 
-    def get_training_tasks_by_token(
+class MockTrainingTaskRepository:
+    """Mock TrainingTaskRepository for testing Phase 4."""
+
+    def __init__(self):
+        self.training_tasks = {}
+        self.training_links = {}
+
+    def get_paginated(
         self,
-        admin_token,
+        admin_token=None,
         status=None,
         limit=20,
         offset=0,
@@ -196,11 +211,22 @@ class MockAdminDB:
         total = len(tasks)
         return tasks[offset:offset+limit], total
 
-    def get_training_task(self, task_id):
+    def get(self, task_id):
         """Get training task by ID."""
         return self.training_tasks.get(str(task_id))
 
-    def get_model_versions(self, status=None, limit=20, offset=0):
+    def get_document_training_tasks(self, document_id):
+        """Get training tasks that used this document."""
+        return self.training_links.get(str(document_id), [])
+
+
+class MockModelVersionRepository:
+    """Mock ModelVersionRepository for testing Phase 4."""
+
+    def __init__(self):
+        self.model_versions = {}
+
+    def get_paginated(self, status=None, limit=20, offset=0):
         """Get model versions with optional filtering."""
         models = list(self.model_versions.values())
         if status:
@@ -214,8 +240,11 @@ 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()
+    mock_training_task_repo = MockTrainingTaskRepository()
+    mock_model_version_repo = MockModelVersionRepository()
 
     # Add test documents
     doc1 = MockAdminDocument(
@@ -231,22 +260,25 @@ def app():
         status="labeled",
     )
 
-    mock_db.documents[str(doc1.document_id)] = doc1
-    mock_db.documents[str(doc2.document_id)] = doc2
-    mock_db.documents[str(doc3.document_id)] = doc3
+    mock_document_repo.documents[str(doc1.document_id)] = doc1
+    mock_document_repo.documents[str(doc2.document_id)] = doc2
+    mock_document_repo.documents[str(doc3.document_id)] = doc3
 
     # Add annotations
-    mock_db.annotations[str(doc1.document_id)] = [
+    mock_annotation_repo.annotations[str(doc1.document_id)] = [
         MockAnnotation(document_id=doc1.document_id, source="manual"),
         MockAnnotation(document_id=doc1.document_id, source="auto"),
     ]
-    mock_db.annotations[str(doc2.document_id)] = [
+    mock_annotation_repo.annotations[str(doc2.document_id)] = [
         MockAnnotation(document_id=doc2.document_id, source="auto"),
         MockAnnotation(document_id=doc2.document_id, source="auto"),
         MockAnnotation(document_id=doc2.document_id, source="auto"),
     ]
     # doc3 has no annotations
 
+    # Share annotation data with document repo for filtering
+    mock_document_repo.annotations = mock_annotation_repo.annotations
+
     # Add training tasks
     task1 = MockTrainingTask(
         name="Training Run 2024-01",
@@ -265,15 +297,18 @@ def app():
         metrics_recall=0.92,
     )
 
-    mock_db.training_tasks[str(task1.task_id)] = task1
-    mock_db.training_tasks[str(task2.task_id)] = task2
+    mock_training_task_repo.training_tasks[str(task1.task_id)] = task1
+    mock_training_task_repo.training_tasks[str(task2.task_id)] = task2
 
     # Add training links (doc1 used in task1)
     link1 = MockTrainingDocumentLink(
         task_id=task1.task_id,
         document_id=doc1.document_id,
     )
-    mock_db.training_links[str(doc1.document_id)] = [link1]
+    mock_training_task_repo.training_links[str(doc1.document_id)] = [link1]
+
+    # Share training links with document repo for filtering
+    mock_document_repo.training_links = mock_training_task_repo.training_links
 
     # Add model versions
     model1 = MockModelVersion(
@@ -296,12 +331,15 @@ def app():
         metrics_recall=0.92,
         document_count=600,
     )
-    mock_db.model_versions[str(model1.version_id)] = model1
-    mock_db.model_versions[str(model2.version_id)] = model2
+    mock_model_version_repo.model_versions[str(model1.version_id)] = model1
+    mock_model_version_repo.model_versions[str(model2.version_id)] = model2
 
     # 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
+    app.dependency_overrides[get_annotation_repository] = lambda: mock_annotation_repo
+    app.dependency_overrides[get_training_task_repository] = lambda: mock_training_task_repo
+    app.dependency_overrides[get_model_version_repository] = lambda: mock_model_version_repo
 
     # Include router
     router = create_training_router()