Files
Imago/tests/test_ai_vision_internals.py
T
bruno 2f2a5e9b4a
CI / Lint & Format (push) Failing after 28s
CI / Tests (push) Has been skipped
CI / Security Scan (push) Failing after 8s
CI / Docker Build (push) Has been skipped
Add admin panel, WebSocket support, and API versioning
Introduce an admin portal (React + Nginx), WebSocket routing, and
API versioning middleware with `/api/v1/` prefix deprecation.
Add master API key authentication, new Prometheus metrics for AI
token consumption and active WebSockets, and extend S3 config
with a public endpoint URL. Update test paths and fixtures to
align with the new routing structure.
2026-06-22 11:25:22 -04:00

96 lines
3.4 KiB
Python

import pytest
import os
import base64
from pathlib import Path
from unittest.mock import patch, MagicMock, AsyncMock
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
@pytest.mark.asyncio
async def test_read_image_variants(tmp_path):
img_file = tmp_path / "test.png"
img_file.write_bytes(b"fake_png_data")
data, mime = await _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", new_callable=AsyncMock)
@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", new_callable=AsyncMock)
@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