diff --git a/penparse/penparse/test_settings.py b/penparse/penparse/test_settings.py index 1eb21f0..634bd3a 100644 --- a/penparse/penparse/test_settings.py +++ b/penparse/penparse/test_settings.py @@ -9,7 +9,7 @@ https://docs.djangoproject.com/en/4.2/topics/settings/ For the full list of settings and their values, see https://docs.djangoproject.com/en/4.2/ref/settings/ """ - +import os from pathlib import Path # Build paths inside the project like this: BASE_DIR / 'subdir'. @@ -130,3 +130,7 @@ DEFAULT_AUTO_FIELD = "django.db.models.BigAutoField" CELERY_BROKER_URL = "amqp://rabbit:rabbit@localhost//" + +OPENAI_API_BASE = os.getenv("OPENAI_API_BASE") +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +OPENAI_MODEL = os.getenv("OPENAI_MODEL", "openai/gpt-4o") diff --git a/penparse/webui/test/test_tasks.py b/penparse/webui/test/test_tasks.py index 0548403..a6488f1 100644 --- a/penparse/webui/test/test_tasks.py +++ b/penparse/webui/test/test_tasks.py @@ -1,104 +1,104 @@ import pytest +import litellm +from unittest.mock import ANY 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 +from ..tasks import process_memo +from ..models import ImageMemo, MemoStatus, User @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 +class TestProcessMemo: - process_memo(sample_image_memo.id) + @pytest.fixture + def sample_image_memo(self, db): + user1 = User.objects.create_user(email="user1@test.com", password="password1") + memo = ImageMemo.objects.create( + author=user1, + status=MemoStatus.Pending, + image_mimetype="image/jpeg", + ) + memo.image.save("test_image.jpg", ContentFile(b"fake image content")) + return memo - 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 == "" + def test_process_memo_success(self, sample_image_memo): + with patch("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) -@pytest.mark.django_db -def test_process_memo_missing_image(sample_image_memo): - default_storage.delete(sample_image_memo.image.name) + 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 == None - 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") + def test_process_memo_missing_image(self, 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 "API Error" in processed_memo.error_message + assert "Image file" in processed_memo.error_message + def test_process_memo_api_error(self, sample_image_memo): + with patch("webui.tasks.litellm") as mock_litellm: + mock_response = MagicMock() + mock_response.choices[0].message = {"content": "Transcribed content"} + mock_litellm.completion.side_effect = litellm.APIError(400, "API Error", "openai", "any") -@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) - 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 - processed_memo = ImageMemo.objects.get(id=sample_image_memo.id) - assert processed_memo.model_name == "test-model" + def test_process_memo_sets_model_name(self, sample_image_memo): + with ( + patch("webui.tasks.litellm") as mock_litellm, + patch("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) -@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" + processed_memo = ImageMemo.objects.get(id=sample_image_memo.id) + assert processed_memo.model_name == "test-model" - process_memo(sample_image_memo.id) + def test_process_memo_uses_correct_api_settings(self, sample_image_memo): + with ( + patch("webui.tasks.litellm") as mock_litellm, + patch("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" - 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, - ) + 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": ANY}, + {"type": "image_url", "image_url": {"url": ANY}}, + ], + } + ] + ), + temperature=0.01, + )