Simplify meme preprocessing and upload failure handling

src/flows/storage/memes.py

@@ -4,7 +4,7 @@
 
 from src.flows.hooks import notify_telegram_on_failure
 from src.storage import ads
-from src.storage.constants import MemeStatus, MemeType
+from src.storage.constants import MemeStatus
 from src.storage.deduplication import (
     deduplicate_pending_meme,
     sweep_file_id_duplicates,
@@ -24,7 +24,7 @@
     download_meme_content_file,
     upload_meme_content_to_tg,
 )
-from src.storage.watermark import add_watermark
+from src.storage.watermark import add_watermark_to_meme_content
 
 
 async def analyse_meme_caption(meme: dict[str, Any]) -> None:
@@ -46,16 +46,6 @@ async def analyse_meme_caption(meme: dict[str, Any]) -> None:
             return
 
 
-async def add_watermark_to_meme_content(meme_content: bytes, meme_type: MemeType) -> bytes | None:
-    if meme_type == MemeType.IMAGE:
-        # we can add watermark only to photos right now
-        watermarked_content = add_watermark(meme_content)
-        if watermarked_content is None:
-            return None
-        return watermarked_content.getvalue()
-    return meme_content
-
-
 async def upload_meme_to_telegram(
     meme: dict[str, Any],
 ) -> dict[str, Any] | None:
@@ -68,9 +58,7 @@ async def upload_meme_to_telegram(
         await update_meme(meme["id"], status=MemeStatus.BROKEN_CONTENT_LINK)
         return None
 
-    watermarked_meme_content = await add_watermark_to_meme_content(
-        meme_original_content, meme["type"]
-    )
+    watermarked_meme_content = add_watermark_to_meme_content(meme_original_content, meme["type"])
     if watermarked_meme_content is None:
         logger.warning(f"Can't add watermark to {meme['id']}/{meme['type']} content")
         return None
@@ -105,15 +93,7 @@ async def _process_unloaded_memes(
             meme = await upload_meme_to_telegram(unloaded_meme)
         except Exception as e:
             logger.warning(f"Meme {unloaded_meme['id']}: upload error: {e}")
-            fail_count += 1
-            consecutive_fails += 1
-            if consecutive_fails >= 5:
-                logger.error(
-                    f"5 consecutive failures — stopping batch. "
-                    f"Processed {i + 1}/{total}, {ok_count} ok, {fail_count} failed."
-                )
-                break
-            continue
+            meme = None
 
         if not meme:
             fail_count += 1

src/storage/watermark.py

@@ -4,6 +4,8 @@
 
 from PIL import Image, ImageDraw, ImageFont, ImageStat
 
+from src.storage.constants import MemeType
+
 
 def draw_text_with_outline(
     draw: ImageDraw,
@@ -113,19 +115,16 @@ def draw_corner_watermark(
         return Image.alpha_composite(base, txt).convert("RGB")
 
 
-# TODO: async?
-def add_watermark(image_content: bytes) -> BytesIO | None:
-    image_bytes = BytesIO(image_content)
+def add_watermark_to_meme_content(meme_content: bytes, meme_type: MemeType) -> bytes | None:
+    if meme_type != MemeType.IMAGE:
+        return meme_content
 
     try:
-        image = draw_corner_watermark(image_bytes, text="@ffmemesbot")
+        image = draw_corner_watermark(BytesIO(meme_content), text="@ffmemesbot")
     except Exception as e:
         print(f"Error while adding watermark: {e}")
         return None
 
     buff = BytesIO()
-    buff.name = "image.jpeg"
     image.save(buff, "JPEG")
-    buff.seek(0)
-
-    return buff
+    return buff.getvalue()

src/tgbot/handlers/upload/moderation.py

@@ -13,10 +13,6 @@
 from src import localizer
 from src.config import settings
 from src.flows.storage.describe_memes import describe_single_meme
-from src.flows.storage.memes import (
-    add_watermark_to_meme_content,
-    upload_meme_content_to_tg,
-)
 from src.observability.sentry import (
     capture_handled_exception,
     capture_handled_issue,
@@ -31,7 +27,8 @@
     resolve_duplicate,
 )
 from src.storage.service import update_meme
-from src.storage.upload import download_meme_content_from_tg
+from src.storage.upload import download_meme_content_from_tg, upload_meme_content_to_tg
+from src.storage.watermark import add_watermark_to_meme_content
 from src.tgbot.handlers.treasury.constants import TrxType
 from src.tgbot.handlers.treasury.payments import pay_if_not_paid_with_alert
 from src.tgbot.handlers.upload.constants import SUPPORTED_LANGUAGES
@@ -273,7 +270,7 @@ async def _uploaded_meme_auto_review(
         )
 
     logging.info(f"Adding watermark to meme {meme['id']} content")
-    watermarked_meme_content = await add_watermark_to_meme_content(image_bytes, meme["type"])
+    watermarked_meme_content = add_watermark_to_meme_content(image_bytes, meme["type"])
     if watermarked_meme_content is None:
         capture_handled_issue(
             "user_upload.watermark_failed",

src/tgbot/handlers/upload/service.py

@@ -14,7 +14,7 @@
     meme_source,
 )
 from src.storage.constants import MemeSourceType, MemeStatus, MemeType
-from src.tgbot.service import get_or_create_meme_source
+from src.tgbot.repo.meme_sources import get_or_create_meme_source
 
 
 async def create_meme_raw_upload(msg: Message) -> dict[str, Any]:
@@ -72,17 +72,6 @@ async def get_user_upload_ban_status(user_id: int) -> dict[str, Any] | None:
     return await fetch_one(query, {"user_id": user_id})
 
 
-async def get_or_create_meme_source_for_meme_upload(
-    meme_upload: dict[str, Any],
-) -> dict[str, Any]:
-    return await get_or_create_meme_source(
-        url="tg://user?id={}".format(meme_upload["user_id"]),
-        type=MemeSourceType.USER_UPLOAD,
-        status=MemeStatus.CREATED,
-        added_by=meme_upload["user_id"],
-    )
-
-
 async def create_meme_from_meme_raw_upload(
     meme_upload: dict[str, Any],
 ) -> dict[str, Any]:
@@ -102,12 +91,18 @@ async def create_meme_from_meme_raw_upload(
     else:
         meme_type = MemeType.IMAGE
 
-    meme_source_for_meme_upload = await get_or_create_meme_source_for_meme_upload(meme_upload)
+    user_id = meme_upload["user_id"]
+    upload_source = await get_or_create_meme_source(
+        url=f"tg://user?id={user_id}",
+        type=MemeSourceType.USER_UPLOAD,
+        status=MemeStatus.CREATED,
+        added_by=user_id,
+    )
 
     query = (
         insert(meme)
         .values(
-            meme_source_id=meme_source_for_meme_upload["id"],
+            meme_source_id=upload_source["id"],
             raw_meme_id=meme_upload["id"],
             status=MemeStatus.CREATED,
             type=meme_type,

tests/storage/test_meme_pipeline_upload.py

@@ -1,19 +1,47 @@
+import logging
 from io import BytesIO
+from unittest.mock import AsyncMock, call
 
 import pytest
+from PIL import Image
 
 from src.flows.storage import memes
+from src.storage import watermark
 from src.storage.constants import MemeType
 
 
-@pytest.mark.asyncio
-async def test_add_watermark_to_meme_content_returns_bytes_for_images(monkeypatch):
-    watermarked_content = BytesIO(b"watermarked-image")
-    watermarked_content.name = "image.jpeg"
+def test_watermark_returns_jpeg_bytes():
+    original = BytesIO()
+    Image.new("RGB", (320, 240), "white").save(original, "PNG")
 
-    monkeypatch.setattr(memes, "add_watermark", lambda content: watermarked_content)
+    result = watermark.add_watermark_to_meme_content(original.getvalue(), MemeType.IMAGE)
 
-    result = await memes.add_watermark_to_meme_content(b"raw-image", MemeType.IMAGE)
+    with Image.open(BytesIO(result)) as image:
+        assert image.format == "JPEG"
+        assert image.size == (320, 240)
+        assert image.getextrema() != ((255, 255),) * 3
 
-    assert result == b"watermarked-image"
-    assert isinstance(result, bytes)
+
+def test_watermark_rejects_invalid_image():
+    assert watermark.add_watermark_to_meme_content(b"invalid", MemeType.IMAGE) is None
+
+
+@pytest.mark.parametrize("meme_type", [MemeType.VIDEO, MemeType.ANIMATION])
+def test_watermark_preserves_non_images(meme_type):
+    content = b"unchanged-media"
+    assert watermark.add_watermark_to_meme_content(content, meme_type) is content
+
+
+async def test_upload_batch_stops_after_five_consecutive_failures(monkeypatch, caplog):
+    error = RuntimeError("Telegram unavailable")
+    unloaded_memes = [{"id": meme_id} for meme_id in range(1, 12)]
+    upload = AsyncMock(
+        side_effect=[None, error, None, error, {"id": 5}, None, error, None, None, error]
+    )
+    monkeypatch.setattr(memes, "upload_meme_to_telegram", upload)
+    monkeypatch.setattr(memes, "get_run_logger", lambda: logging.getLogger(__name__))
+
+    await memes._process_unloaded_memes(unloaded_memes, "Telegram")
+
+    assert upload.await_args_list == [call(meme) for meme in unloaded_memes[:10]]
+    assert caplog.text.count("5 consecutive failures") == 1

tests/tgbot/test_upload_meme.py

@@ -1,7 +1,11 @@
 from unittest.mock import AsyncMock
 
+from sqlalchemy import delete, select
 from telegram.error import BadRequest
+from tests.factories import FIXED_DT, create_user
 
+from src.database import engine, fetch_one, meme_source, user
+from src.tgbot.handlers.upload.service import create_meme_from_meme_raw_upload
 from src.tgbot.handlers.upload.upload_meme import _edit_upload_callback_caption
 
 
@@ -20,3 +24,33 @@ def __init__(self) -> None:
     await _edit_upload_callback_caption(message, "caption", reply_markup=None)
 
     message.edit_caption.assert_awaited_once()
+
+
+async def test_uploads_reuse_the_uploaders_source():
+    user_id = 910_000_000_013
+    async with engine.begin() as conn:
+        await create_user(conn, user_id)
+
+    upload = {
+        "user_id": user_id,
+        "forward_origin": None,
+        "date": FIXED_DT,
+        "media": {"file_id": "uploaded-photo"},
+        "language_code": "ru",
+    }
+    try:
+        first = await create_meme_from_meme_raw_upload({**upload, "id": 1})
+        second = await create_meme_from_meme_raw_upload({**upload, "id": 2})
+
+        source = await fetch_one(
+            select(meme_source).where(meme_source.c.id == first["meme_source_id"])
+        )
+        assert first["meme_source_id"] == second["meme_source_id"]
+        assert source["url"] == f"tg://user?id={user_id}"
+        assert source["type"] == "user upload"
+        assert source["status"] == "created"
+        assert source["added_by"] == user_id
+    finally:
+        async with engine.begin() as conn:
+            await conn.execute(delete(meme_source).where(meme_source.c.added_by == user_id))
+            await conn.execute(delete(user).where(user.c.id == user_id))

tests/tgbot/test_upload_moderation.py

@@ -198,7 +198,7 @@ async def test_auto_review_does_not_revive_exact_file_id_duplicate(monkeypatch):
         patch.object(
             moderation,
             "add_watermark_to_meme_content",
-            new=AsyncMock(return_value=b"watermarked"),
+            return_value=b"watermarked",
         ),
         patch.object(
             moderation,