From 3412a5ccaa0d019976c84a2d307ffc9907954c93 Mon Sep 17 00:00:00 2001 From: zwitschi Date: Sun, 10 May 2026 14:24:07 +0200 Subject: [PATCH] feat: enhance message handling with content pattern matching and update environment configuration --- .env.example | 9 ++++ main.py | 30 ++++++++--- maintenance.py | 55 ++++++++++++++++++-- tests/test_main_scheduler.py | 98 ++++++++++++++++++++++++++++++++++++ tests/test_maintenance.py | 38 ++++++++++++++ 5 files changed, 221 insertions(+), 9 deletions(-) diff --git a/.env.example b/.env.example index 09759a0..cdab87b 100644 --- a/.env.example +++ b/.env.example @@ -1,2 +1,11 @@ # Replace with your actual Discord webhook URL DISCORD_WEBHOOK_URL=https://discord.com/api/webhooks// + +# Replace with your Discord bot token +DISCORD_BOT_TOKEN= + +# Replace with your Discord channel ID +DISCORD_CHANNEL_ID= + +# Replace with your Discord guild/server ID +DISCORD_GUILD_ID= diff --git a/main.py b/main.py index a6100ec..679b602 100644 --- a/main.py +++ b/main.py @@ -50,6 +50,7 @@ logging.basicConfig( load_dotenv() WEBHOOK_URL = os.getenv('DISCORD_WEBHOOK_URL') +TEST_MESSAGE_DELETE_PATTERN = r"test notification" def get_state() -> dict: @@ -158,13 +159,16 @@ def create_embed(type: str, tz_list: list[str] | None = None) -> dict: return embed -def send_notification(message: str) -> None: +def send_notification(message: str) -> bool: """ Send a notification to the Discord webhook. + + Returns: + bool: True when the webhook accepted the notification, False otherwise. """ if not WEBHOOK_URL: logging.error("WEBHOOK_URL not set") - return + return False _update_state( last_type=message, @@ -192,6 +196,7 @@ def send_notification(message: str) -> None: logging.info(f"Notification sent: {message}") _update_state(last_success_at=datetime.now(), last_status_code=response.status_code) + return True else: logging.error( f"Failed to send notification: {response.status_code} - " @@ -199,9 +204,23 @@ def send_notification(message: str) -> None: ) _update_state(last_status_code=response.status_code, last_error=response.text) + return False except requests.RequestException as e: logging.error(f"Error sending notification: {e}") _update_state(last_error=str(e)) + return False + + +def _schedule_startup_test_cleanup(test_sent: bool) -> None: + """Schedule one-time cleanup for the startup test notification.""" + if not test_sent: + return + + def cleanup_startup_test_message() -> schedule.CancelJob: + delete_old_messages(1, content_pattern=TEST_MESSAGE_DELETE_PATTERN) + return schedule.CancelJob + + schedule.every(1).minutes.do(cleanup_startup_test_message) def schedule_notification(interval: str, at: str, type: str) -> None: @@ -237,10 +256,9 @@ def main() -> None: logging.info("Scheduler started.") - # Test the notification on startup - send_notification("test") - # delete the test message after a short delay to keep the channel clean - schedule.every(1).minutes.do(delete_old_messages, 1) + # Send one startup test message and cleanup only if send succeeded. + test_sent = send_notification("test") + _schedule_startup_test_cleanup(test_sent) # delete old messages on startup to clean up any previous notifications # delete_old_messages(6) diff --git a/maintenance.py b/maintenance.py index 863f7d2..2e680b5 100644 --- a/maintenance.py +++ b/maintenance.py @@ -1,5 +1,6 @@ import logging import os +import re import time from datetime import datetime, timedelta @@ -42,15 +43,55 @@ def should_delete_message( webhook_id: str, author_id: str, cutoff: int, + content_pattern: str | None = None, ) -> bool: message_timestamp = int(parse_message_timestamp(message).timestamp()) return ( message_timestamp <= cutoff and message.get("webhook_id") == webhook_id and message.get("author", {}).get("id") == author_id + and message_matches_pattern(message, content_pattern) ) +def message_matches_pattern(message: dict, content_pattern: str | None = None) -> bool: + """Return True when message content/embed text matches the optional pattern.""" + if not content_pattern: + return True + + text_chunks: list[str] = [] + content = message.get("content") + if isinstance(content, str) and content: + text_chunks.append(content) + + embeds = message.get("embeds") + if isinstance(embeds, list): + for embed in embeds: + if not isinstance(embed, dict): + continue + title = embed.get("title") + description = embed.get("description") + if isinstance(title, str) and title: + text_chunks.append(title) + if isinstance(description, str) and description: + text_chunks.append(description) + + footer = embed.get("footer") + if isinstance(footer, dict): + footer_text = footer.get("text") + if isinstance(footer_text, str) and footer_text: + text_chunks.append(footer_text) + + if not text_chunks: + return False + + searchable_text = "\n".join(text_chunks) + try: + return re.search(content_pattern, searchable_text, flags=re.IGNORECASE) is not None + except re.error: + return content_pattern.lower() in searchable_text.lower() + + def get_rate_limit_retry_after(response: requests.Response) -> float | None: header_retry_after = parse_float(response.headers.get("Retry-After")) if header_retry_after is not None: @@ -137,7 +178,7 @@ def find_last_message_by_author( return None -def fetch_messages_to_delete(headers: dict, channel_id: str, webhook_id: str, author_id: str, cutoff: int, last_message_id: str | None = None) -> tuple[list[dict], str | None]: +def fetch_messages_to_delete(headers: dict, channel_id: str, webhook_id: str, author_id: str, cutoff: int, last_message_id: str | None = None, content_pattern: str | None = None) -> tuple[list[dict], str | None]: """ Fetch messages from the channel that are older than the cutoff timestamp and sent by the webhook. Uses pagination with the 'before' parameter to resume from the last processed message. @@ -180,6 +221,7 @@ def fetch_messages_to_delete(headers: dict, channel_id: str, webhook_id: str, au webhook_id, author_id, cutoff, + content_pattern, ): delete_list.append(build_delete_entry(message)) @@ -240,7 +282,7 @@ def delete_message(headers: dict, channel_id: str, message_id: str) -> tuple[boo return False, None, False -def delete_old_messages(minutes: int = 6) -> None: +def delete_old_messages(minutes: int = 6, content_pattern: str | None = None) -> None: """ Delete all messages sent by the webhook in the last `minutes` minutes. Uses a dynamic slowdown to avoid hitting Discord API rate limits and pagination to fetch all messages. @@ -283,6 +325,7 @@ def delete_old_messages(minutes: int = 6) -> None: webhook_id, author_id, cutoff, + content_pattern, ): anchor_message = build_delete_entry(last_author_message) deleted, wait_seconds, abort_batch = delete_message( @@ -303,7 +346,13 @@ def delete_old_messages(minutes: int = 6) -> None: while True: delete_list, next_last_message_id = fetch_messages_to_delete( - headers, discord_channel_id, webhook_id, author_id, cutoff, last_message_id + headers, + discord_channel_id, + webhook_id, + author_id, + cutoff, + last_message_id, + content_pattern, ) if not delete_list: diff --git a/tests/test_main_scheduler.py b/tests/test_main_scheduler.py index 8c67ba0..cfb7c14 100644 --- a/tests/test_main_scheduler.py +++ b/tests/test_main_scheduler.py @@ -23,3 +23,101 @@ def test_get_next_scheduled_event(): nxt = main.get_next_scheduled_event(now) assert nxt["type"] == "reminder" assert nxt["at"].hour == 11 and nxt["at"].minute == 15 + + +def test_schedule_startup_test_cleanup_when_sent(monkeypatch): + captured: dict[str, object] = {} + delete_calls: list[tuple[int, str | None]] = [] + + class FakeEvery: + @property + def minutes(self): + return self + + def do(self, fn, *args, **kwargs): + captured["job"] = lambda: fn(*args, **kwargs) + return object() + + monkeypatch.setattr(main.schedule, "every", lambda n: FakeEvery()) + monkeypatch.setattr( + main, + "delete_old_messages", + lambda minutes, content_pattern=None: delete_calls.append( + (minutes, content_pattern)), + ) + + main._schedule_startup_test_cleanup(True) + + assert "job" in captured + result = captured["job"]() + assert result == main.schedule.CancelJob + assert delete_calls == [(1, main.TEST_MESSAGE_DELETE_PATTERN)] + + +def test_schedule_startup_test_cleanup_skips_when_not_sent(monkeypatch): + called = {"value": False} + + class FakeEvery: + @property + def minutes(self): + return self + + def do(self, fn, *args, **kwargs): + called["value"] = True + return object() + + monkeypatch.setattr(main.schedule, "every", lambda n: FakeEvery()) + + main._schedule_startup_test_cleanup(False) + + assert called["value"] is False + + +def test_main_sends_startup_test_and_deletes_it(monkeypatch): + send_calls: list[str] = [] + delete_calls: list[tuple[int, str | None]] = [] + scheduled_jobs: dict[int, list[object]] = {1: [], 5: []} + + monkeypatch.setattr(main, "start_dashboard", lambda: None) + monkeypatch.setattr(main, "schedule_notification", + lambda interval, at, type: None) + + class FakeEvery: + def __init__(self, minutes_value: int): + self.minutes_value = minutes_value + + @property + def minutes(self): + return self + + def do(self, fn, *args, **kwargs): + scheduled_jobs.setdefault(self.minutes_value, []).append( + lambda: fn(*args, **kwargs) + ) + return object() + + monkeypatch.setattr(main.schedule, "every", lambda n: FakeEvery(n)) + + def fake_send_notification(message: str) -> bool: + send_calls.append(message) + return True + + def fake_delete_old_messages(minutes: int = 6, content_pattern: str | None = None): + delete_calls.append((minutes, content_pattern)) + + def fake_run_pending(): + for job in scheduled_jobs.get(1, []): + job() + raise KeyboardInterrupt() + + monkeypatch.setattr(main, "send_notification", fake_send_notification) + monkeypatch.setattr(main, "delete_old_messages", fake_delete_old_messages) + monkeypatch.setattr(main.schedule, "run_pending", fake_run_pending) + monkeypatch.setattr(main.time, "sleep", lambda s: None) + monkeypatch.setenv("DISCORD_WEBHOOK_URL", "http://example.com/webhook") + main.WEBHOOK_URL = "http://example.com/webhook" + + main.main() + + assert send_calls == ["test"] + assert delete_calls == [(1, main.TEST_MESSAGE_DELETE_PATTERN)] diff --git a/tests/test_maintenance.py b/tests/test_maintenance.py index 8bcce99..b088760 100644 --- a/tests/test_maintenance.py +++ b/tests/test_maintenance.py @@ -53,6 +53,44 @@ def test_should_delete_message(): ) +def test_message_matches_pattern_content_and_embeds(): + message = { + "content": "This is a smoke test payload", + "embeds": [ + {"title": "Reminder", "description": "Half-time in 5 minutes"} + ], + } + + assert maintenance.message_matches_pattern(message, r"smoke test") + assert maintenance.message_matches_pattern(message, r"half-time") + assert not maintenance.message_matches_pattern(message, r"does-not-match") + + +def test_should_delete_message_with_content_pattern(): + ts = int(datetime(2026, 1, 1, 10, 0, 0, tzinfo=timezone.utc).timestamp()) + message = { + "timestamp": "2026-01-01T10:00:00Z", + "webhook_id": "w", + "author": {"id": "a"}, + "embeds": [{"description": "This is a test notification."}], + } + + assert maintenance.should_delete_message( + message, + webhook_id="w", + author_id="a", + cutoff=ts, + content_pattern=r"test notification", + ) + assert not maintenance.should_delete_message( + message, + webhook_id="w", + author_id="a", + cutoff=ts, + content_pattern=r"production-only", + ) + + def test_get_rate_limit_retry_after_header_priority(): response = DummyResponse( headers={