275 lines
9.5 KiB
Python
275 lines
9.5 KiB
Python
"""
|
|
Tests for the AsyncProcessingService class.
|
|
"""
|
|
|
|
import tempfile
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from inference.data.async_request_db import AsyncRequest
|
|
from inference.web.workers.async_queue import AsyncTask, AsyncTaskQueue
|
|
from inference.web.services.async_processing import AsyncProcessingService, AsyncSubmitResult
|
|
from inference.web.config import AsyncConfig, StorageConfig
|
|
from inference.web.rate_limiter import RateLimiter
|
|
|
|
|
|
@pytest.fixture
|
|
def async_service(mock_db, mock_inference_service, rate_limiter, storage_config):
|
|
"""Create an AsyncProcessingService for testing."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
async_config = 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,
|
|
)
|
|
|
|
queue = AsyncTaskQueue(max_size=10, worker_count=1)
|
|
|
|
service = AsyncProcessingService(
|
|
inference_service=mock_inference_service,
|
|
db=mock_db,
|
|
queue=queue,
|
|
rate_limiter=rate_limiter,
|
|
async_config=async_config,
|
|
storage_config=storage_config,
|
|
)
|
|
|
|
yield service
|
|
|
|
# Cleanup
|
|
if service._queue._started:
|
|
service.stop()
|
|
|
|
|
|
class TestAsyncProcessingService:
|
|
"""Tests for AsyncProcessingService."""
|
|
|
|
def test_submit_request_success(self, async_service, mock_db):
|
|
"""Test successful request submission."""
|
|
mock_db.create_request.return_value = "test-request-id"
|
|
|
|
result = async_service.submit_request(
|
|
api_key="test-api-key",
|
|
file_content=b"fake pdf content",
|
|
filename="test.pdf",
|
|
content_type="application/pdf",
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.request_id is not None
|
|
assert result.estimated_wait_seconds >= 0
|
|
assert result.error is None
|
|
|
|
def test_submit_request_creates_db_record(self, async_service, mock_db):
|
|
"""Test that submission creates database record."""
|
|
async_service.submit_request(
|
|
api_key="test-api-key",
|
|
file_content=b"fake pdf content",
|
|
filename="test.pdf",
|
|
content_type="application/pdf",
|
|
)
|
|
|
|
mock_db.create_request.assert_called_once()
|
|
call_kwargs = mock_db.create_request.call_args[1]
|
|
assert call_kwargs["api_key"] == "test-api-key"
|
|
assert call_kwargs["filename"] == "test.pdf"
|
|
assert call_kwargs["content_type"] == "application/pdf"
|
|
|
|
def test_submit_request_saves_file(self, async_service, mock_db):
|
|
"""Test that submission saves file to temp directory."""
|
|
content = b"fake pdf content"
|
|
|
|
result = async_service.submit_request(
|
|
api_key="test-api-key",
|
|
file_content=content,
|
|
filename="test.pdf",
|
|
content_type="application/pdf",
|
|
)
|
|
|
|
# File should exist in temp directory
|
|
temp_dir = async_service._async_config.temp_upload_dir
|
|
files = list(temp_dir.iterdir())
|
|
|
|
# Note: file may be cleaned up quickly if queue processes it
|
|
# So we just check that the operation succeeded
|
|
assert result.success is True
|
|
|
|
def test_submit_request_records_rate_limit(self, async_service, mock_db, rate_limiter):
|
|
"""Test that submission records rate limit event."""
|
|
async_service.submit_request(
|
|
api_key="test-api-key",
|
|
file_content=b"fake pdf content",
|
|
filename="test.pdf",
|
|
content_type="application/pdf",
|
|
)
|
|
|
|
# Rate limiter should have recorded the request
|
|
mock_db.record_rate_limit_event.assert_called()
|
|
|
|
def test_start_and_stop(self, async_service):
|
|
"""Test starting and stopping the service."""
|
|
async_service.start()
|
|
|
|
assert async_service._queue._started is True
|
|
assert async_service._cleanup_thread is not None
|
|
assert async_service._cleanup_thread.is_alive()
|
|
|
|
async_service.stop()
|
|
|
|
assert async_service._queue._started is False
|
|
|
|
def test_process_task_success(self, async_service, mock_db, mock_inference_service, sample_task):
|
|
"""Test successful task processing."""
|
|
async_service._process_task(sample_task)
|
|
|
|
# Should update status to processing
|
|
mock_db.update_status.assert_called_with(sample_task.request_id, "processing")
|
|
|
|
# Should complete the request
|
|
mock_db.complete_request.assert_called_once()
|
|
call_kwargs = mock_db.complete_request.call_args[1]
|
|
assert call_kwargs["request_id"] == sample_task.request_id
|
|
assert "document_id" in call_kwargs
|
|
|
|
def test_process_task_pdf(self, async_service, mock_db, mock_inference_service, sample_task):
|
|
"""Test processing a PDF task."""
|
|
async_service._process_task(sample_task)
|
|
|
|
# Should call process_pdf for .pdf files
|
|
mock_inference_service.process_pdf.assert_called_once()
|
|
|
|
def test_process_task_image(self, async_service, mock_db, mock_inference_service):
|
|
"""Test processing an image task."""
|
|
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
|
f.write(b"fake image content")
|
|
task = AsyncTask(
|
|
request_id="image-task",
|
|
api_key="test-api-key",
|
|
file_path=Path(f.name),
|
|
filename="test.png",
|
|
)
|
|
|
|
async_service._process_task(task)
|
|
|
|
# Should call process_image for image files
|
|
mock_inference_service.process_image.assert_called_once()
|
|
|
|
def test_process_task_failure(self, async_service, mock_db, mock_inference_service, sample_task):
|
|
"""Test task processing failure."""
|
|
mock_inference_service.process_pdf.side_effect = Exception("Processing failed")
|
|
|
|
async_service._process_task(sample_task)
|
|
|
|
# Should update status to failed
|
|
mock_db.update_status.assert_called()
|
|
last_call = mock_db.update_status.call_args_list[-1]
|
|
assert last_call[0][1] == "failed" # status
|
|
assert "Processing failed" in last_call[1]["error_message"]
|
|
|
|
def test_process_task_file_not_found(self, async_service, mock_db):
|
|
"""Test task processing with missing file."""
|
|
task = AsyncTask(
|
|
request_id="missing-file-task",
|
|
api_key="test-api-key",
|
|
file_path=Path("/nonexistent/file.pdf"),
|
|
filename="test.pdf",
|
|
)
|
|
|
|
async_service._process_task(task)
|
|
|
|
# Should fail with file not found
|
|
mock_db.update_status.assert_called()
|
|
last_call = mock_db.update_status.call_args_list[-1]
|
|
assert last_call[0][1] == "failed"
|
|
|
|
def test_process_task_cleans_up_file(self, async_service, mock_db, mock_inference_service):
|
|
"""Test that task processing cleans up the uploaded file."""
|
|
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
|
|
f.write(b"fake pdf content")
|
|
file_path = Path(f.name)
|
|
|
|
task = AsyncTask(
|
|
request_id="cleanup-task",
|
|
api_key="test-api-key",
|
|
file_path=file_path,
|
|
filename="test.pdf",
|
|
)
|
|
|
|
async_service._process_task(task)
|
|
|
|
# File should be deleted
|
|
assert not file_path.exists()
|
|
|
|
def test_estimate_wait(self, async_service):
|
|
"""Test wait time estimation."""
|
|
# Empty queue
|
|
wait = async_service._estimate_wait()
|
|
assert wait == 0
|
|
|
|
def test_cleanup_orphan_files(self, async_service, mock_db):
|
|
"""Test cleanup of orphan files."""
|
|
# Create the async upload directory
|
|
temp_dir = async_service._async_config.temp_upload_dir
|
|
temp_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
orphan_file = temp_dir / "orphan-request.pdf"
|
|
orphan_file.write_bytes(b"orphan content")
|
|
|
|
# Set file mtime to old
|
|
import os
|
|
old_time = time.time() - 7200
|
|
os.utime(orphan_file, (old_time, old_time))
|
|
|
|
# Mock database to say file doesn't exist
|
|
mock_db.get_request.return_value = None
|
|
|
|
# Mock the storage helper to return the same directory as the fixture
|
|
with patch("inference.web.services.async_processing.get_storage_helper") as mock_storage:
|
|
mock_helper = MagicMock()
|
|
mock_helper.get_uploads_base_path.return_value = temp_dir
|
|
mock_storage.return_value = mock_helper
|
|
|
|
count = async_service._cleanup_orphan_files()
|
|
|
|
assert count == 1
|
|
assert not orphan_file.exists()
|
|
|
|
def test_save_upload(self, async_service):
|
|
"""Test saving uploaded file."""
|
|
content = b"test content"
|
|
|
|
file_path = async_service._save_upload(
|
|
request_id="test-save",
|
|
filename="test.pdf",
|
|
content=content,
|
|
)
|
|
|
|
assert file_path.exists()
|
|
assert file_path.read_bytes() == content
|
|
assert file_path.suffix == ".pdf"
|
|
|
|
# Cleanup
|
|
file_path.unlink()
|
|
|
|
def test_save_upload_preserves_extension(self, async_service):
|
|
"""Test that save_upload preserves file extension."""
|
|
content = b"test content"
|
|
|
|
# Test various extensions
|
|
for ext in [".pdf", ".png", ".jpg", ".jpeg"]:
|
|
file_path = async_service._save_upload(
|
|
request_id=f"test-{ext}",
|
|
filename=f"test{ext}",
|
|
content=content,
|
|
)
|
|
|
|
assert file_path.suffix == ext
|
|
file_path.unlink()
|