133 lines
3.6 KiB
Python
133 lines
3.6 KiB
Python
"""
|
|
Test fixtures for web API tests.
|
|
"""
|
|
|
|
import tempfile
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
|
|
from inference.data.async_request_db import ApiKeyConfig, AsyncRequestDB
|
|
from inference.data.models import AsyncRequest
|
|
from inference.web.workers.async_queue import AsyncTask, AsyncTaskQueue
|
|
from inference.web.services.async_processing import AsyncProcessingService
|
|
from inference.web.config import AsyncConfig, StorageConfig
|
|
from inference.web.core.rate_limiter import RateLimiter
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_db():
|
|
"""Create a mock AsyncRequestDB."""
|
|
db = MagicMock(spec=AsyncRequestDB)
|
|
|
|
# Default return values
|
|
db.is_valid_api_key.return_value = True
|
|
db.get_api_key_config.return_value = ApiKeyConfig(
|
|
api_key="test-api-key",
|
|
name="Test Key",
|
|
is_active=True,
|
|
requests_per_minute=10,
|
|
max_concurrent_jobs=3,
|
|
max_file_size_mb=50,
|
|
)
|
|
db.count_active_jobs.return_value = 0
|
|
db.get_queue_position.return_value = 1
|
|
|
|
return db
|
|
|
|
|
|
@pytest.fixture
|
|
def rate_limiter(mock_db):
|
|
"""Create a RateLimiter with mock database."""
|
|
return RateLimiter(mock_db)
|
|
|
|
|
|
@pytest.fixture
|
|
def task_queue():
|
|
"""Create an AsyncTaskQueue."""
|
|
return AsyncTaskQueue(max_size=10, worker_count=1)
|
|
|
|
|
|
@pytest.fixture
|
|
def async_config():
|
|
"""Create an AsyncConfig for testing."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
yield AsyncConfig(
|
|
queue_max_size=10,
|
|
worker_count=1,
|
|
task_timeout_seconds=30,
|
|
result_retention_days=7,
|
|
temp_upload_dir=Path(tmpdir) / "async",
|
|
max_file_size_mb=10,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def storage_config():
|
|
"""Create a StorageConfig for testing."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
yield StorageConfig(
|
|
upload_dir=Path(tmpdir) / "uploads",
|
|
result_dir=Path(tmpdir) / "results",
|
|
max_file_size_mb=50,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_inference_service():
|
|
"""Create a mock InferenceService."""
|
|
service = MagicMock()
|
|
service.is_initialized = True
|
|
service.gpu_available = False
|
|
|
|
# Mock process_pdf to return a successful result
|
|
mock_result = MagicMock()
|
|
mock_result.document_id = "test-doc"
|
|
mock_result.success = True
|
|
mock_result.document_type = "invoice"
|
|
mock_result.fields = {"InvoiceNumber": "12345", "Amount": "1000.00"}
|
|
mock_result.confidence = {"InvoiceNumber": 0.95, "Amount": 0.92}
|
|
mock_result.detections = []
|
|
mock_result.errors = []
|
|
mock_result.visualization_path = None
|
|
|
|
service.process_pdf.return_value = mock_result
|
|
service.process_image.return_value = mock_result
|
|
|
|
return service
|
|
|
|
|
|
# Valid UUID for testing
|
|
TEST_REQUEST_UUID = "550e8400-e29b-41d4-a716-446655440000"
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_async_request():
|
|
"""Create a sample AsyncRequest."""
|
|
return AsyncRequest(
|
|
request_id=UUID(TEST_REQUEST_UUID),
|
|
api_key="test-api-key",
|
|
status="pending",
|
|
filename="test.pdf",
|
|
file_size=1024,
|
|
content_type="application/pdf",
|
|
expires_at=datetime.utcnow() + timedelta(days=7),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_task():
|
|
"""Create a sample AsyncTask."""
|
|
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
|
|
f.write(b"fake pdf content")
|
|
return AsyncTask(
|
|
request_id=TEST_REQUEST_UUID,
|
|
api_key="test-api-key",
|
|
file_path=Path(f.name),
|
|
filename="test.pdf",
|
|
created_at=datetime.utcnow(),
|
|
)
|