- Implement tests for database generator to ensure proper session handling. - Create tests for EXIF extraction and conversion functions. - Add tests for image-related endpoints, ensuring proper data retrieval and isolation between clients. - Develop tests for OCR functionality, including language detection and text extraction. - Introduce tests for the image processing pipeline, covering success and failure scenarios. - Validate rate limiting functionality and ensure independent counters for different clients. - Implement scraper tests to verify HTML content fetching and error handling. - Add unit tests for various services, including storage and filename generation. - Establish worker entry point for ARQ to handle background image processing tasks.
107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
import pytest
|
|
import uuid
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from app.models.image import Image, ProcessingStatus
|
|
from app.services.pipeline import process_image_pipeline
|
|
from unittest.mock import patch, AsyncMock
|
|
|
|
def create_test_image(client_id, **kwargs):
|
|
u = str(uuid.uuid4())
|
|
defaults = {
|
|
"uuid": u,
|
|
"client_id": client_id,
|
|
"original_name": "test.jpg",
|
|
"filename": f"{u}.jpg",
|
|
"file_path": f"fake/{u}.jpg",
|
|
}
|
|
defaults.update(kwargs)
|
|
return Image(**defaults)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_full_success(db_session: AsyncSession, client_a):
|
|
image = create_test_image(client_id=client_a.id)
|
|
db_session.add(image)
|
|
await db_session.commit()
|
|
await db_session.refresh(image)
|
|
|
|
mock_exif = {"make": "Canon", "model": "EOS R"}
|
|
mock_ocr = {"text": "Hello", "has_text": True, "confidence": 0.9}
|
|
mock_ai = {"description": "A photo", "tags": ["tag1"], "confidence": 0.8, "model": "gemini"}
|
|
|
|
with patch("app.services.pipeline.extract_exif", return_value=mock_exif), \
|
|
patch("app.services.pipeline.extract_text", return_value=mock_ocr), \
|
|
patch("app.services.pipeline.analyze_image", return_value=mock_ai):
|
|
|
|
await process_image_pipeline(image.id, db_session)
|
|
|
|
await db_session.refresh(image)
|
|
assert image.processing_status == ProcessingStatus.DONE
|
|
assert image.exif_make == "Canon"
|
|
assert image.ocr_text == "Hello"
|
|
assert image.ai_description == "A photo"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_partial_failure(db_session: AsyncSession, client_a):
|
|
image = create_test_image(client_id=client_a.id)
|
|
db_session.add(image)
|
|
await db_session.commit()
|
|
await db_session.refresh(image)
|
|
|
|
mock_ocr = {"text": "Hello", "has_text": True}
|
|
mock_ai = {"description": "A photo", "tags": ["tag1"]}
|
|
|
|
# Step 1 fails, Step 2 & 3 succeed
|
|
with patch("app.services.pipeline.extract_exif", side_effect=Exception("EXIF error")), \
|
|
patch("app.services.pipeline.extract_text", return_value=mock_ocr), \
|
|
patch("app.services.pipeline.analyze_image", return_value=mock_ai):
|
|
|
|
await process_image_pipeline(image.id, db_session)
|
|
|
|
await db_session.refresh(image)
|
|
assert image.processing_status == ProcessingStatus.DONE # Still DONE because AI succeeded
|
|
assert "EXIF error" in image.processing_error
|
|
assert image.ocr_text == "Hello"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_total_failure(db_session: AsyncSession, client_a):
|
|
image = create_test_image(client_id=client_a.id)
|
|
db_session.add(image)
|
|
await db_session.commit()
|
|
await db_session.refresh(image)
|
|
|
|
# All steps fail
|
|
with patch("app.services.pipeline.extract_exif", side_effect=Exception("EXIF error")), \
|
|
patch("app.services.pipeline.extract_text", side_effect=Exception("OCR error")), \
|
|
patch("app.services.pipeline.analyze_image", side_effect=Exception("AI error")):
|
|
|
|
await process_image_pipeline(image.id, db_session)
|
|
|
|
await db_session.refresh(image)
|
|
assert image.processing_status == ProcessingStatus.ERROR
|
|
assert "EXIF error" in image.processing_error
|
|
assert "OCR error" in image.processing_error
|
|
assert "AI error" in image.processing_error
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_ocr_fallback(db_session: AsyncSession, client_a):
|
|
image = create_test_image(client_id=client_a.id)
|
|
db_session.add(image)
|
|
await db_session.commit()
|
|
await db_session.refresh(image)
|
|
|
|
mock_exif = {}
|
|
mock_ocr_empty = {"has_text": False}
|
|
mock_ai_ocr = {"text": "AI extracted text", "has_text": True}
|
|
mock_ai_vision = {"description": "A photo"}
|
|
|
|
with patch("app.services.pipeline.extract_exif", return_value=mock_exif), \
|
|
patch("app.services.pipeline.extract_text", return_value=mock_ocr_empty), \
|
|
patch("app.services.pipeline.extract_text_with_ai", return_value=mock_ai_ocr), \
|
|
patch("app.services.pipeline.analyze_image", return_value=mock_ai_vision):
|
|
|
|
await process_image_pipeline(image.id, db_session)
|
|
|
|
await db_session.refresh(image)
|
|
assert image.ocr_text == "AI extracted text"
|
|
assert image.ocr_has_text is True
|