WIP

tests/web/test_augmentation_routes.py

@@ -5,7 +5,75 @@ TDD Phase 5: RED - Write tests first, then implement to pass.
 """
 
 import pytest
+from unittest.mock import MagicMock, patch
+from fastapi import FastAPI
 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
+
+
+TEST_ADMIN_TOKEN = "test-admin-token-12345"
+TEST_DOCUMENT_UUID = "550e8400-e29b-41d4-a716-446655440001"
+TEST_DATASET_UUID = "660e8400-e29b-41d4-a716-446655440001"
+
+
+@pytest.fixture
+def admin_token() -> str:
+    """Provide admin token for testing."""
+    return TEST_ADMIN_TOKEN
+
+
+@pytest.fixture
+def mock_admin_db() -> MagicMock:
+    """Create a mock AdminDB 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)
+    return mock
+
+
+@pytest.fixture
+def admin_client(mock_admin_db: MagicMock) -> TestClient:
+    """Create test client with admin authentication."""
+    app = FastAPI()
+
+    # Override dependencies
+    def get_token_override():
+        return TEST_ADMIN_TOKEN
+
+    def get_db_override():
+        return mock_admin_db
+
+    app.dependency_overrides[validate_admin_token] = get_token_override
+    app.dependency_overrides[get_admin_db] = get_db_override
+
+    # Include router - the router already has /augmentation prefix
+    # so we add /api/v1/admin to get /api/v1/admin/augmentation
+    router = create_augmentation_router()
+    app.include_router(router, prefix="/api/v1/admin")
+
+    return TestClient(app)
+
+
+@pytest.fixture
+def unauthenticated_client(mock_admin_db: 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
+
+    app.dependency_overrides[get_admin_db] = get_db_override
+
+    router = create_augmentation_router()
+    app.include_router(router, prefix="/api/v1/admin")
+
+    return TestClient(app)
 
 
 class TestAugmentationTypesEndpoint:
@@ -34,10 +102,10 @@ class TestAugmentationTypesEndpoint:
         assert "stage" in aug_type
 
     def test_list_augmentation_types_unauthorized(
-        self, admin_client: TestClient
+        self, unauthenticated_client: TestClient
     ) -> None:
         """Test that unauthorized request is rejected."""
-        response = admin_client.get("/api/v1/admin/augmentation/types")
+        response = unauthenticated_client.get("/api/v1/admin/augmentation/types")
 
         assert response.status_code == 401
 
@@ -74,16 +142,30 @@ class TestAugmentationPreviewEndpoint:
         admin_client: TestClient,
         admin_token: str,
         sample_document_id: str,
+        mock_admin_db: MagicMock,
     ) -> None:
         """Test previewing augmentation on a document."""
-        response = admin_client.post(
-            f"/api/v1/admin/augmentation/preview/{sample_document_id}",
-            headers={"X-Admin-Token": admin_token},
-            json={
-                "augmentation_type": "gaussian_noise",
-                "params": {"std": 15},
-            },
-        )
+        # Mock document exists
+        mock_document = MagicMock()
+        mock_document.images_dir = "/fake/path"
+        mock_admin_db.get_document.return_value = mock_document
+
+        # Create a fake image (100x100 RGB)
+        fake_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+
+        with patch(
+            "inference.web.services.augmentation_service.AugmentationService._load_document_page"
+        ) as mock_load:
+            mock_load.return_value = fake_image
+
+            response = admin_client.post(
+                f"/api/v1/admin/augmentation/preview/{sample_document_id}",
+                headers={"X-Admin-Token": admin_token},
+                json={
+                    "augmentation_type": "gaussian_noise",
+                    "params": {"std": 15},
+                },
+            )
 
         assert response.status_code == 200
         data = response.json()
@@ -136,18 +218,32 @@ class TestAugmentationPreviewConfigEndpoint:
         admin_client: TestClient,
         admin_token: str,
         sample_document_id: str,
+        mock_admin_db: MagicMock,
     ) -> None:
         """Test previewing full config on a document."""
-        response = admin_client.post(
-            f"/api/v1/admin/augmentation/preview-config/{sample_document_id}",
-            headers={"X-Admin-Token": admin_token},
-            json={
-                "gaussian_noise": {"enabled": True, "probability": 1.0},
-                "lighting_variation": {"enabled": True, "probability": 1.0},
-                "preserve_bboxes": True,
-                "seed": 42,
-            },
-        )
+        # Mock document exists
+        mock_document = MagicMock()
+        mock_document.images_dir = "/fake/path"
+        mock_admin_db.get_document.return_value = mock_document
+
+        # Create a fake image (100x100 RGB)
+        fake_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+
+        with patch(
+            "inference.web.services.augmentation_service.AugmentationService._load_document_page"
+        ) as mock_load:
+            mock_load.return_value = fake_image
+
+            response = admin_client.post(
+                f"/api/v1/admin/augmentation/preview-config/{sample_document_id}",
+                headers={"X-Admin-Token": admin_token},
+                json={
+                    "gaussian_noise": {"enabled": True, "probability": 1.0},
+                    "lighting_variation": {"enabled": True, "probability": 1.0},
+                    "preserve_bboxes": True,
+                    "seed": 42,
+                },
+            )
 
         assert response.status_code == 200
         data = response.json()
@@ -164,8 +260,14 @@ class TestAugmentationBatchEndpoint:
         admin_client: TestClient,
         admin_token: str,
         sample_dataset_id: str,
+        mock_admin_db: 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
+
         response = admin_client.post(
             "/api/v1/admin/augmentation/batch",
             headers={"X-Admin-Token": admin_token},
@@ -250,12 +352,10 @@ class TestAugmentedDatasetsListEndpoint:
 @pytest.fixture
 def sample_document_id() -> str:
     """Provide a sample document ID for testing."""
-    # This would need to be created in test setup
-    return "test-document-id"
+    return TEST_DOCUMENT_UUID
 
 
 @pytest.fixture
 def sample_dataset_id() -> str:
     """Provide a sample dataset ID for testing."""
-    # This would need to be created in test setup
-    return "test-dataset-id"
+    return TEST_DATASET_UUID