implement tests for worker

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")

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 ..tasks import process_memo
-from penparse.webui.models import ImageMemo, MemoStatus
+from ..models import ImageMemo, MemoStatus, User
-
-
-@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):
+class TestProcessMemo:
-    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)
+    @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)
+    def test_process_memo_success(self, sample_image_memo):
-        assert processed_memo.status == MemoStatus.Done
+        with patch("webui.tasks.litellm") as mock_litellm:
-        assert processed_memo.content == "Transcribed content"
+            mock_response = MagicMock()
-        assert processed_memo.error_message == ""
+            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
+            processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
-def test_process_memo_missing_image(sample_image_memo):
+            assert processed_memo.status == MemoStatus.Done
-    default_storage.delete(sample_image_memo.image.name)
+            assert processed_memo.content == "Transcribed content"
+            assert processed_memo.error_message == None
 
-    process_memo(sample_image_memo.id)
+    def test_process_memo_missing_image(self, 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.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
+        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
+            process_memo(sample_image_memo.id)
-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.status == MemoStatus.Error
+            assert "API Error" in processed_memo.error_message
 
-        processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
+    def test_process_memo_sets_model_name(self, sample_image_memo):
-        assert processed_memo.model_name == "test-model"
+        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
+            processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
-def test_process_memo_uses_correct_api_settings(sample_image_memo):
+            assert processed_memo.model_name == "test-model"
-    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)
+    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"
+            process_memo(sample_image_memo.id)
-        assert mock_litellm.api_key == "test-api-key"
+
-        mock_litellm.completion.assert_called_once_with(
+            assert mock_litellm.api_base == "https://test-api-base.com"
-            model="test-model",
+            assert mock_litellm.api_key == "test-api-key"
-            messages=pytest.approx(
+            mock_litellm.completion.assert_called_once_with(
-                [
+                model="test-model",
-                    {
+                messages=pytest.approx(
-                        "role": "user",
+                    [
-                        "content": [
+                        {
-                            {"type": "text", "text": pytest.ANY},
+                            "role": "user",
-                            {"type": "image_url", "image_url": {"url": pytest.ANY}},
+                            "content": [
-                        ],
+                                {"type": "text", "text": ANY},
-                    }
+                                {"type": "image_url", "image_url": {"url": ANY}},
-                ]
+                            ],
-            ),
+                        }
-            temperature=0.01,
+                    ]
-        )
+                ),
+                temperature=0.01,
+            )