From dba3388567818cb5c60be9d406d3e1cccc84b3ea Mon Sep 17 00:00:00 2001 From: salmonik Date: Wed, 29 Apr 2026 14:22:03 +0200 Subject: Refactor database connections and add integration tests - Use context managers for sqlite3 connections in cron, vclogs, and database utils - Add URL helper utility for yt-dlp with clean_url_for_ytdlp function - Add slow integration test markers and new integration test files - Fix test imports to use conftest directly instead of tests.conftest - Improve downloader error handling and use COPYPARTY_URL_UPLOAD env var - Update pytest.ini with slow test markers and warning filters - Document slow test usage in AGENTS.md --- cogs/cron.py | 11 ++------- cogs/downloader.py | 14 ++++++++++- cogs/music/media_processing.py | 2 ++ cogs/music/play_commands.py | 2 ++ cogs/utils/database.py | 17 ++++++++++++- cogs/utils/url_helper.py | 21 ++++++++++++++++ cogs/vclogs.py | 55 ++++++++++++++++++++---------------------- cogs/web.py | 2 ++ 8 files changed, 84 insertions(+), 40 deletions(-) create mode 100644 cogs/utils/url_helper.py (limited to 'cogs') diff --git a/cogs/cron.py b/cogs/cron.py index b4f494a..aad43d2 100644 --- a/cogs/cron.py +++ b/cogs/cron.py @@ -30,13 +30,10 @@ class Cron(commands.Cog): def load_tasks(self): try: - conn = get_connection() - try: + with get_connection() as conn: c = conn.cursor() c.execute("SELECT id, cron_expression, type, channel_id, active, audio_filename, message FROM cron_jobs") rows = c.fetchall() - finally: - conn.close() self.scheduled_tasks = [ { @@ -57,8 +54,7 @@ class Cron(commands.Cog): def save_tasks(self): try: - conn = get_connection() - try: + with get_connection() as conn: c = conn.cursor() c.execute("PRAGMA foreign_keys = ON") @@ -78,10 +74,7 @@ class Cron(commands.Cog): task.get("message"), ), ) - conn.commit() - finally: - conn.close() except Exception as e: logger.error(f"Error saving cron jobs: {e}") diff --git a/cogs/downloader.py b/cogs/downloader.py index d829ed7..900bafb 100644 --- a/cogs/downloader.py +++ b/cogs/downloader.py @@ -12,6 +12,7 @@ import re import yt_dlp from discord.ext import commands import aiofiles +from cogs.utils.url_helper import clean_url_for_ytdlp logger = logging.getLogger(__name__) @@ -61,6 +62,10 @@ class Downloader(commands.Cog): file_data = response["content"] await ctx.respond(file=discord.File(file_data)) else: + error_msg = response["content"] + if 'details' in response: + error_msg += f"\n{response['details']}" + await ctx.respond(f"Błąd: {error_msg}", ephemeral=True) logger.error(f"Error: {response['content']}") if 'details' in response: logger.error(f"Details: {response['details']}") @@ -68,6 +73,7 @@ class Downloader(commands.Cog): @staticmethod async def download_standalone(url: str, catbox: bool = False, fformat: str = "mp4", filename_user: str | None = None): try: + url = clean_url_for_ytdlp(url) logger.info(f"Downloading {url} as {fformat}") with tempfile.TemporaryDirectory(delete=False) as temp_dir: @@ -157,12 +163,18 @@ class Downloader(commands.Cog): files = {'f': (filename, content, 'application/octet-stream')} headers = {'Cookie': f'cppwd={os.getenv("COPYPARTY_PASSWORD")}'} params = {'j': '1', 'ck': 'b2'} - base_url = "https://catbox.2137697.xyz/sperma/" + base_url = os.getenv("COPYPARTY_URL_UPLOAD") response = requests.post(base_url, files=files, headers=headers, params=params) if response.status_code in [200, 201]: return Downloader._handle_upload_success(response, filename, temp_dir) + elif response.status_code == 413: + return Downloader._create_error_response( + "File too large for upload", + "Plik jest za duży, by wysłać go na serwer (limit to ok. 100MB)", + temp_dir + ) else: return Downloader._create_error_response( f"Server returned {response.status_code}", diff --git a/cogs/music/media_processing.py b/cogs/music/media_processing.py index b9b814c..91bd406 100644 --- a/cogs/music/media_processing.py +++ b/cogs/music/media_processing.py @@ -9,6 +9,7 @@ import os import yt_dlp from youtubesearchpython import VideosSearch +from cogs.utils.url_helper import clean_url_for_ytdlp logger = logging.getLogger(__name__) @@ -36,6 +37,7 @@ class MediaDownloader: } try: + url = clean_url_for_ytdlp(url) with yt_dlp.YoutubeDL(ytdl_opts) as ydl: info_dict = ydl.extract_info(url, download=True) # Get the base filename without extension diff --git a/cogs/music/play_commands.py b/cogs/music/play_commands.py index c53dbc7..bdb021d 100644 --- a/cogs/music/play_commands.py +++ b/cogs/music/play_commands.py @@ -19,6 +19,7 @@ from discord.ext import commands from cogs.music import audio_playback, media_processing from cogs.utils.file_search import UniversalSearcher from cogs.utils.settings import get_uploadurl +from cogs.utils.url_helper import clean_url_for_ytdlp logger = logging.getLogger(__name__) @@ -193,6 +194,7 @@ class PlayCommands(commands.Cog): } loop = asyncio.get_event_loop() + url = clean_url_for_ytdlp(url) with YoutubeDL(ytdl_opts) as ydl: info_dict = await loop.run_in_executor( None, lambda: ydl.extract_info(url, download=True) diff --git a/cogs/utils/database.py b/cogs/utils/database.py index 67cbc9f..ea7d356 100644 --- a/cogs/utils/database.py +++ b/cogs/utils/database.py @@ -1,6 +1,7 @@ import os import sqlite3 import logging +from contextlib import contextmanager logger = logging.getLogger(__name__) @@ -10,7 +11,8 @@ _initialized = False def get_connection(): - return sqlite3.connect(get_db_path()) + conn = sqlite3.connect(get_db_path()) + return _ConnectionContextManager(conn) def get_db_path(): @@ -22,6 +24,19 @@ def set_db_path(path): DB_PATH = path +class _ConnectionContextManager: + """Context manager for sqlite3 connections that auto-closes.""" + def __init__(self, conn): + self.conn = conn + + def __enter__(self): + return self.conn + + def __exit__(self, exc_type, exc_val, exc_tb): + self.conn.close() + return False + + def init_database(): global _initialized if _initialized: diff --git a/cogs/utils/url_helper.py b/cogs/utils/url_helper.py new file mode 100644 index 0000000..eac06f1 --- /dev/null +++ b/cogs/utils/url_helper.py @@ -0,0 +1,21 @@ +# This file is part of Senator Sperma +# Licensed under the GPL-3.0 License +# Copyright (C) 2026 Salmonik @ salmonik@tuta.io +# https://codeberg.org/salmonik/senator-sperma + + +def clean_url_for_ytdlp(url: str) -> str: + """Clean URL by removing everything after the first '&'. + + This prevents yt-dlp from interpreting additional query parameters + (like &list=...) as playlist URLs. + + Example: + Input: https://www.youtube.com/watch?v=lPGipwoJiOM&list=RDlPGipwoJiOM& + Output: https://www.youtube.com/watch?v=lPGipwoJiOM + """ + # Find the first '&' and truncate + first_ampersand = url.find('&') + if first_ampersand != -1: + return url[:first_ampersand] + return url diff --git a/cogs/vclogs.py b/cogs/vclogs.py index 404ffe0..e5ca202 100644 --- a/cogs/vclogs.py +++ b/cogs/vclogs.py @@ -23,11 +23,10 @@ init_database() def load_activity_data(): try: - conn = get_connection() - c = conn.cursor() - c.execute("SELECT user_id, join_time, total_active, total_afk, is_afk, afk_start, active_start, channel_name FROM voice_activity") - rows = c.fetchall() - conn.close() + with get_connection() as conn: + c = conn.cursor() + c.execute("SELECT user_id, join_time, total_active, total_afk, is_afk, afk_start, active_start, channel_name FROM voice_activity") + rows = c.fetchall() return { row[0]: { @@ -48,30 +47,28 @@ def load_activity_data(): def save_activity_data(data): try: - conn = get_connection() - c = conn.cursor() - c.execute("PRAGMA foreign_keys = ON") - - c.execute("DELETE FROM voice_activity") - - for user_id, user_data in data.items(): - c.execute( - """INSERT INTO voice_activity (user_id, join_time, total_active, total_afk, is_afk, afk_start, active_start, channel_name) - VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", - ( - user_id, - user_data.get("join_time"), - user_data.get("total_active", 0), - user_data.get("total_afk", 0), - int(user_data.get("is_afk", False)), - user_data.get("afk_start"), - user_data.get("active_start"), - user_data.get("channel_name"), - ), - ) - - conn.commit() - conn.close() + with get_connection() as conn: + c = conn.cursor() + c.execute("PRAGMA foreign_keys = ON") + + c.execute("DELETE FROM voice_activity") + + for user_id, user_data in data.items(): + c.execute( + """INSERT INTO voice_activity (user_id, join_time, total_active, total_afk, is_afk, afk_start, active_start, channel_name) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", + ( + user_id, + user_data.get("join_time"), + user_data.get("total_active", 0), + user_data.get("total_afk", 0), + int(user_data.get("is_afk", False)), + user_data.get("afk_start"), + user_data.get("active_start"), + user_data.get("channel_name"), + ), + ) + conn.commit() logger.info(f"Saved voice activity for {len(data)} users") except Exception as e: logger.error(f"Error saving voice activity: {e}") diff --git a/cogs/web.py b/cogs/web.py index eea13b1..fd22f9f 100644 --- a/cogs/web.py +++ b/cogs/web.py @@ -11,6 +11,7 @@ from bs4 import BeautifulSoup import urllib.parse import yt_dlp import logging +from cogs.utils.url_helper import clean_url_for_ytdlp logger = logging.getLogger(__name__) @@ -24,6 +25,7 @@ class Web(commands.Cog): if "youtube.com/watch" not in url: return None + url = clean_url_for_ytdlp(url) logger.info("youtube link") try: ydl_opts = {"quiet": True, "nowarnings": True} -- cgit v1.2.3