From 0b14ec20fb0c18b49b3d0574768921cf6dfc5bbb Mon Sep 17 00:00:00 2001 From: Daniil Okhlopkov <5613295+ohld@users.noreply.github.com> Date: Mon, 9 Nov 2026 08:30:58 -0300 Subject: [PATCH 1/4] refactor: move shared watermark policy into storage --- src/flows/storage/memes.py | 18 +++--------------- src/storage/watermark.py | 12 +++++++++++- src/tgbot/handlers/upload/moderation.py | 9 +++------ tests/storage/test_meme_pipeline_upload.py | 11 ++++------- tests/tgbot/test_upload_moderation.py | 2 +- 5 files changed, 22 insertions(+), 30 deletions(-) diff --git a/src/flows/storage/memes.py b/src/flows/storage/memes.py index 92503f63..334afa8a 100644 --- a/src/flows/storage/memes.py +++ b/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 diff --git a/src/storage/watermark.py b/src/storage/watermark.py index 9c1c7f53..997a97df 100644 --- a/src/storage/watermark.py +++ b/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,7 +115,6 @@ 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) @@ -129,3 +130,12 @@ def add_watermark(image_content: bytes) -> BytesIO | None: buff.seek(0) return buff + + +def add_watermark_to_meme_content(meme_content: bytes, meme_type: MemeType) -> bytes | None: + if meme_type != MemeType.IMAGE: + return meme_content + watermarked_content = add_watermark(meme_content) + if watermarked_content is None: + return None + return watermarked_content.getvalue() diff --git a/src/tgbot/handlers/upload/moderation.py b/src/tgbot/handlers/upload/moderation.py index 37b69a6c..421750bd 100644 --- a/src/tgbot/handlers/upload/moderation.py +++ b/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", diff --git a/tests/storage/test_meme_pipeline_upload.py b/tests/storage/test_meme_pipeline_upload.py index 455cd8ac..3bf7f1d0 100644 --- a/tests/storage/test_meme_pipeline_upload.py +++ b/tests/storage/test_meme_pipeline_upload.py @@ -1,19 +1,16 @@ from io import BytesIO -import pytest - -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): +def test_add_watermark_to_meme_content_returns_bytes_for_images(monkeypatch): watermarked_content = BytesIO(b"watermarked-image") watermarked_content.name = "image.jpeg" - monkeypatch.setattr(memes, "add_watermark", lambda content: watermarked_content) + monkeypatch.setattr(watermark, "add_watermark", lambda content: watermarked_content) - result = await memes.add_watermark_to_meme_content(b"raw-image", MemeType.IMAGE) + result = watermark.add_watermark_to_meme_content(b"raw-image", MemeType.IMAGE) assert result == b"watermarked-image" assert isinstance(result, bytes) diff --git a/tests/tgbot/test_upload_moderation.py b/tests/tgbot/test_upload_moderation.py index 63949cdb..0653daaa 100644 --- a/tests/tgbot/test_upload_moderation.py +++ b/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, From 8f987940bd336a26807c9d039dccc0f397ef690b Mon Sep 17 00:00:00 2001 From: Daniil Okhlopkov <5613295+ohld@users.noreply.github.com> Date: Mon, 9 Nov 2026 08:35:26 -0300 Subject: [PATCH 2/4] refactor: unify ingestion batch failure handling --- src/flows/storage/memes.py | 10 +--------- tests/storage/test_meme_pipeline_upload.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/src/flows/storage/memes.py b/src/flows/storage/memes.py index 334afa8a..1dfc4222 100644 --- a/src/flows/storage/memes.py +++ b/src/flows/storage/memes.py @@ -93,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 diff --git a/tests/storage/test_meme_pipeline_upload.py b/tests/storage/test_meme_pipeline_upload.py index 3bf7f1d0..fed7b931 100644 --- a/tests/storage/test_meme_pipeline_upload.py +++ b/tests/storage/test_meme_pipeline_upload.py @@ -1,5 +1,8 @@ +import logging from io import BytesIO +from unittest.mock import AsyncMock, call +from src.flows.storage import memes from src.storage import watermark from src.storage.constants import MemeType @@ -14,3 +17,18 @@ def test_add_watermark_to_meme_content_returns_bytes_for_images(monkeypatch): assert result == b"watermarked-image" assert isinstance(result, bytes) + + +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 From 7e00073ab1a583b9ee15c1369da1a04886cec578 Mon Sep 17 00:00:00 2001 From: Daniil Okhlopkov <5613295+ohld@users.noreply.github.com> Date: Mon, 9 Nov 2026 08:38:25 -0300 Subject: [PATCH 3/4] refactor: inline single-use upload source wrapper --- src/tgbot/handlers/upload/service.py | 21 +++++++---------- tests/tgbot/test_upload_meme.py | 34 ++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 13 deletions(-) diff --git a/src/tgbot/handlers/upload/service.py b/src/tgbot/handlers/upload/service.py index cadc260f..dfc1e00f 100644 --- a/src/tgbot/handlers/upload/service.py +++ b/src/tgbot/handlers/upload/service.py @@ -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, diff --git a/tests/tgbot/test_upload_meme.py b/tests/tgbot/test_upload_meme.py index c35db637..05344aa0 100644 --- a/tests/tgbot/test_upload_meme.py +++ b/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)) From 89321c6c60681385bc93061a2f8d97ddd3e5fc15 Mon Sep 17 00:00:00 2001 From: Daniil Okhlopkov <5613295+ohld@users.noreply.github.com> Date: Mon, 9 Nov 2026 10:34:03 -0300 Subject: [PATCH 4/4] refactor: collapse watermark buffer layer and test media contracts --- src/storage/watermark.py | 21 ++++------------- src/tgbot/handlers/upload/service.py | 2 +- tests/storage/test_meme_pipeline_upload.py | 27 ++++++++++++++++------ 3 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/storage/watermark.py b/src/storage/watermark.py index 997a97df..db4c35e7 100644 --- a/src/storage/watermark.py +++ b/src/storage/watermark.py @@ -115,27 +115,16 @@ def draw_corner_watermark( return Image.alpha_composite(base, txt).convert("RGB") -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 - - -def add_watermark_to_meme_content(meme_content: bytes, meme_type: MemeType) -> bytes | None: - if meme_type != MemeType.IMAGE: - return meme_content - watermarked_content = add_watermark(meme_content) - if watermarked_content is None: - return None - return watermarked_content.getvalue() + return buff.getvalue() diff --git a/src/tgbot/handlers/upload/service.py b/src/tgbot/handlers/upload/service.py index dfc1e00f..fc86a547 100644 --- a/src/tgbot/handlers/upload/service.py +++ b/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]: diff --git a/tests/storage/test_meme_pipeline_upload.py b/tests/storage/test_meme_pipeline_upload.py index fed7b931..077de52a 100644 --- a/tests/storage/test_meme_pipeline_upload.py +++ b/tests/storage/test_meme_pipeline_upload.py @@ -2,21 +2,34 @@ 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 -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") + + result = watermark.add_watermark_to_meme_content(original.getvalue(), MemeType.IMAGE) + + with Image.open(BytesIO(result)) as image: + assert image.format == "JPEG" + assert image.size == (320, 240) + assert image.getextrema() != ((255, 255),) * 3 + - monkeypatch.setattr(watermark, "add_watermark", lambda content: watermarked_content) +def test_watermark_rejects_invalid_image(): + assert watermark.add_watermark_to_meme_content(b"invalid", MemeType.IMAGE) is None - result = watermark.add_watermark_to_meme_content(b"raw-image", MemeType.IMAGE) - assert result == b"watermarked-image" - assert isinstance(result, bytes) +@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):