Bruno Charest
cc99fea20a
- 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.
95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
import pytest
|
|
import os
|
|
import base64
|
|
from pathlib import Path
|
|
from unittest.mock import patch, MagicMock
|
|
from app.services.ai_vision import (
|
|
_read_image, _extract_json, _usage_tokens_gemini,
|
|
_generate_gemini, _generate_openrouter, _generate,
|
|
extract_text_with_ai, analyze_image
|
|
)
|
|
from app.config import settings
|
|
|
|
def test_read_image_variants(tmp_path):
|
|
img_file = tmp_path / "test.png"
|
|
img_file.write_bytes(b"fake_png_data")
|
|
|
|
data, mime = _read_image(str(img_file))
|
|
assert data == b"fake_png_data"
|
|
assert mime == "image/png"
|
|
|
|
def test_extract_json_bad_cases():
|
|
assert _extract_json(None) is None
|
|
assert _extract_json("No JSON here") is None
|
|
assert _extract_json("Broken { json : 1") is None
|
|
|
|
def test_usage_tokens_gemini_none():
|
|
assert _usage_tokens_gemini(None) == (None, None)
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.usage_metadata = None
|
|
assert _usage_tokens_gemini(mock_resp) == (None, None)
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("app.services.ai_vision._get_client")
|
|
@patch("app.services.ai_vision.settings")
|
|
async def test_generate_gemini_no_key(mock_settings, mock_get_client):
|
|
mock_settings.GEMINI_API_KEY = None
|
|
res = await _generate_gemini("test prompt")
|
|
assert res["text"] is None
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("httpx.AsyncClient.post")
|
|
@patch("app.services.ai_vision.settings")
|
|
async def test_generate_openrouter_success(mock_settings, mock_post):
|
|
mock_settings.OPENROUTER_API_KEY = "test_key"
|
|
mock_settings.OPENROUTER_MODEL = "test_model"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"choices": [{"message": {"content": "openrouter test response"}}],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5}
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
res = await _generate_openrouter("test prompt")
|
|
assert res["text"] == "openrouter test response"
|
|
assert res["usage"] == (10, 5)
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("app.services.ai_vision._generate_openrouter")
|
|
@patch("app.services.ai_vision.settings")
|
|
async def test_generate_dispatcher_openrouter(mock_settings, mock_gen_or):
|
|
mock_settings.AI_PROVIDER = "openrouter"
|
|
await _generate("test")
|
|
mock_gen_or.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("app.services.ai_vision._generate")
|
|
@patch("app.services.ai_vision._read_image")
|
|
@patch("app.services.ai_vision.settings")
|
|
async def test_extract_text_with_ai(mock_settings, mock_read_image, mock_generate):
|
|
mock_settings.AI_ENABLED = True
|
|
mock_settings.AI_PROVIDER = "gemini"
|
|
mock_read_image.return_value = (b"fake", "image/jpeg")
|
|
mock_generate.return_value = {
|
|
"text": '{"text": "Extracted text", "language": "en", "confidence": 0.9}'
|
|
}
|
|
|
|
res = await extract_text_with_ai("fake/path.jpg")
|
|
assert res["text"] == "Extracted text"
|
|
assert res["language"] == "en"
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("app.services.ai_vision._generate")
|
|
@patch("app.services.ai_vision._read_image")
|
|
@patch("app.services.ai_vision.settings")
|
|
async def test_analyze_image_failure(mock_settings, mock_read_image, mock_generate):
|
|
mock_settings.AI_ENABLED = True
|
|
mock_read_image.return_value = (b"fake", "image/jpeg")
|
|
mock_generate.return_value = {"text": "Invalid JSON", "usage": (None, None)}
|
|
|
|
res = await analyze_image("fake/path.jpg")
|
|
assert res["description"] is None
|