105 lines
3.7 KiB
Python
105 lines
3.7 KiB
Python
import pytest
|
|
from django.core.files.base import ContentFile
|
|
from django.core.files.storage import default_storage
|
|
from unittest.mock import patch, MagicMock
|
|
from penparse.webui.tasks import process_memo
|
|
from penparse.webui.models import ImageMemo, MemoStatus
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_image_memo(db):
|
|
memo = ImageMemo.objects.create(
|
|
status=MemoStatus.Pending, image_mimetype="image/jpeg"
|
|
)
|
|
memo.image.save("test_image.jpg", ContentFile(b"fake image content"))
|
|
return memo
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_process_memo_success(sample_image_memo):
|
|
with patch("penparse.webui.tasks.litellm") as mock_litellm:
|
|
mock_response = MagicMock()
|
|
mock_response.choices[0].message = {"content": "Transcribed content"}
|
|
mock_litellm.completion.return_value = mock_response
|
|
|
|
process_memo(sample_image_memo.id)
|
|
|
|
processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
|
|
assert processed_memo.status == MemoStatus.Done
|
|
assert processed_memo.content == "Transcribed content"
|
|
assert processed_memo.error_message == ""
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_process_memo_missing_image(sample_image_memo):
|
|
default_storage.delete(sample_image_memo.image.name)
|
|
|
|
process_memo(sample_image_memo.id)
|
|
|
|
processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
|
|
assert processed_memo.status == MemoStatus.Error
|
|
assert "Image file" in processed_memo.error_message
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_process_memo_api_error(sample_image_memo):
|
|
with patch("penparse.webui.tasks.litellm") as mock_litellm:
|
|
mock_litellm.completion.side_effect = mock_litellm.APIError("API Error")
|
|
|
|
process_memo(sample_image_memo.id)
|
|
|
|
processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
|
|
assert processed_memo.status == MemoStatus.Error
|
|
assert "API Error" in processed_memo.error_message
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_process_memo_sets_model_name(sample_image_memo):
|
|
with (
|
|
patch("penparse.webui.tasks.litellm") as mock_litellm,
|
|
patch("penparse.webui.tasks.settings") as mock_settings,
|
|
):
|
|
mock_response = MagicMock()
|
|
mock_response.choices[0].message = {"content": "Transcribed content"}
|
|
mock_litellm.completion.return_value = mock_response
|
|
mock_settings.OPENAI_MODEL = "test-model"
|
|
|
|
process_memo(sample_image_memo.id)
|
|
|
|
processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
|
|
assert processed_memo.model_name == "test-model"
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_process_memo_uses_correct_api_settings(sample_image_memo):
|
|
with (
|
|
patch("penparse.webui.tasks.litellm") as mock_litellm,
|
|
patch("penparse.webui.tasks.settings") as mock_settings,
|
|
):
|
|
mock_response = MagicMock()
|
|
mock_response.choices[0].message = {"content": "Transcribed content"}
|
|
mock_litellm.completion.return_value = mock_response
|
|
mock_settings.OPENAI_API_BASE = "https://test-api-base.com"
|
|
mock_settings.OPENAI_API_KEY = "test-api-key"
|
|
mock_settings.OPENAI_MODEL = "test-model"
|
|
|
|
process_memo(sample_image_memo.id)
|
|
|
|
assert mock_litellm.api_base == "https://test-api-base.com"
|
|
assert mock_litellm.api_key == "test-api-key"
|
|
mock_litellm.completion.assert_called_once_with(
|
|
model="test-model",
|
|
messages=pytest.approx(
|
|
[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": pytest.ANY},
|
|
{"type": "image_url", "image_url": {"url": pytest.ANY}},
|
|
],
|
|
}
|
|
]
|
|
),
|
|
temperature=0.01,
|
|
)
|