""" 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 an orphan file temp_dir = async_service._async_config.temp_upload_dir 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 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()