WIP

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)