Imago/tests/test_pipeline_extended.py
Bruno Charest cc99fea20a
Some checks failed
CI / Lint & Format (push) Has been cancelled
CI / Tests (push) Has been cancelled
CI / Security Scan (push) Has been cancelled
CI / Docker Build (push) Has been cancelled
Add comprehensive test suite for image processing and related services
- 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.
2026-02-24 11:22:10 -05:00

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