fix: addressing feedback from the PR

Author: felipemontoya

backend/openedx_ai_badges/edxapp_wrapper/contentstore.py

@@ -2,22 +2,27 @@
 Contentstore module generalized definitions.
 """
 
+from functools import lru_cache
 from importlib import import_module
 
 from django.conf import settings
 
 
+@lru_cache(maxsize=1)
+def _get_backend():
+    """Return the configured contentstore backend module."""
+    return import_module(settings.OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND)
+
+
 def get_static_content():
     """
     Wrapper for `xmodule.contentstore.content.StaticContent` in edx-platform.
     """
-    backend = import_module(settings.OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND)
-    return backend.StaticContent
+    return _get_backend().StaticContent
 
 
 def update_course_run_asset(*args, **kwargs):
     """
     Wrapper for `cms.djangoapps.contentstore.views.assets.update_course_run_asset`.
     """
-    backend = import_module(settings.OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND)
-    return backend.update_course_run_asset(*args, **kwargs)
+    return _get_backend().update_course_run_asset(*args, **kwargs)

backend/openedx_ai_badges/processors/badge_image_upload_processor.py

@@ -41,11 +41,18 @@ def upload_image_to_assets(self, course_id, b64_string, badge_id):
             # Strip data URL prefix if present (e.g. "data:image/png;base64,...")
             if "," in b64_string:
                 b64_string = b64_string.split(",", 1)[1]
-            image_bytes = base64.b64decode(b64_string)
+            image_bytes = base64.b64decode(b64_string, validate=True)
         except Exception:  # pylint: disable=broad-except
             logger.exception("Failed to decode base64 badge image for badge %s", badge_id)
             return None
 
+        max_size = getattr(settings, "OPENEDX_AI_BADGES_MAX_IMAGE_SIZE_BYTES", 5 * 1024 * 1024)
+        if len(image_bytes) > max_size:
+            logger.error(
+                "Badge image for badge %s exceeds maximum allowed size (%d bytes)", badge_id, max_size
+            )
+            return None
+
         filename = f"badge_{badge_id}.png"
         file_obj = InMemoryUploadedFile(
             file=io.BytesIO(image_bytes),

backend/openedx_ai_badges/settings/common.py

@@ -30,10 +30,13 @@ def plugin_settings(settings):
     if badges_workflow_dir not in settings.WORKFLOW_TEMPLATE_DIRS:
         settings.WORKFLOW_TEMPLATE_DIRS.append(badges_workflow_dir)
 
-    # Wrappers for openedx-platform functions
+    # -------------------------
+    # Contentstore wrapper
+    # -------------------------
     settings.OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND = (
         "openedx_ai_badges.edxapp_wrapper.backends.contentstore_r_v1"
     )
+    settings.OPENEDX_AI_BADGES_MAX_IMAGE_SIZE_BYTES = 1 * 1024 * 1024  # 1 MB
 
     # -------------------------
     # MIT DCC Badge API

backend/openedx_ai_badges/settings/production.py

@@ -31,3 +31,14 @@ def plugin_settings(settings):
         settings.MIT_DCC_BADGE_API_HEALTH_URL = settings.ENV_TOKENS.get(
             "MIT_DCC_BADGE_API_HEALTH_URL", settings.MIT_DCC_BADGE_API_HEALTH_URL
         )
+
+    # -------------------------
+    # Contentstore wrapper
+    # -------------------------
+    if hasattr(settings, "ENV_TOKENS"):
+        settings.OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND = settings.ENV_TOKENS.get(
+            "OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND", settings.OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND
+        )
+        settings.OPENEDX_AI_BADGES_MAX_IMAGE_SIZE_BYTES = settings.ENV_TOKENS.get(
+            "OPENEDX_AI_BADGES_MAX_IMAGE_SIZE_BYTES", settings.OPENEDX_AI_BADGES_MAX_IMAGE_SIZE_BYTES
+        )

backend/tests/test_badge_image_upload_processor.py

@@ -78,9 +78,9 @@ def test_skips_conversion_when_already_course_key(self, processor, settings):
 
         with patch("openedx_ai_badges.processors.badge_image_upload_processor.update_course_run_asset", mock_upload), \
              patch("openedx_ai_badges.processors.badge_image_upload_processor.get_static_content", mock_get_static), \
-             patch("openedx_ai_badges.processors.badge_image_upload_processor.CourseKey") as mock_ck:
+             patch.object(CourseKey, "from_string") as mock_from_string:
             processor.upload_image_to_assets(real_key, MINIMAL_PNG_B64, BADGE_ID)
-            mock_ck.from_string.assert_not_called()
+            mock_from_string.assert_not_called()
 
         assert mock_upload.call_args[0][0] is real_key
 
@@ -106,6 +106,18 @@ def test_returns_none_on_upload_failure(self, processor, settings):
 
         assert url is None
 
+    def test_returns_none_when_image_exceeds_max_size(self, processor, settings):
+        settings.LMS_ROOT_URL = LMS_ROOT
+        settings.OPENEDX_AI_BADGES_MAX_IMAGE_SIZE_BYTES = 1  # 1 byte limit
+        mock_upload, mock_get_static = _make_upload_mock()
+
+        with patch("openedx_ai_badges.processors.badge_image_upload_processor.update_course_run_asset", mock_upload), \
+             patch("openedx_ai_badges.processors.badge_image_upload_processor.get_static_content", mock_get_static):
+            url = processor.upload_image_to_assets(COURSE_ID_STR, MINIMAL_PNG_B64, BADGE_ID)
+
+        assert url is None
+        mock_upload.assert_not_called()
+
     def test_uploaded_file_has_correct_name_and_content_type(self, processor, settings):
         settings.LMS_ROOT_URL = LMS_ROOT
         mock_upload, mock_get_static = _make_upload_mock()

backend/tests/test_contentstore_wrapper.py

@@ -8,6 +8,17 @@
 
 from unittest.mock import MagicMock, patch
 
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def clear_backend_cache():
+    """Clear the lru_cache on _get_backend between tests."""
+    from openedx_ai_badges.edxapp_wrapper.contentstore import _get_backend
+    _get_backend.cache_clear()
+    yield
+    _get_backend.cache_clear()
+
 
 class TestGetStaticContent:
 
@@ -34,6 +45,18 @@ def test_uses_configured_backend_setting(self, settings):
 
         mock_import.assert_called_once_with("my.custom.backend")
 
+    def test_backend_is_cached_across_calls(self, settings):
+        fake_backend = MagicMock()
+        fake_backend.StaticContent = MagicMock()
+        settings.OPENEDX_AI_BADGES_CONTENTSTORE_BACKEND = "fake.backend"
+
+        with patch("openedx_ai_badges.edxapp_wrapper.contentstore.import_module", return_value=fake_backend) as mock_import:
+            from openedx_ai_badges.edxapp_wrapper.contentstore import get_static_content
+            get_static_content()
+            get_static_content()
+
+        mock_import.assert_called_once()
+
 
 class TestUpdateCourseRunAsset: