1 files changed, 44 insertions, 41 deletions

diff --git a/tests/cogs/test_cron.py b/tests/cogs/test_cron.py
index 91fcb82..28ccfe1 100644
--- a/tests/cogs/test_cron.py
+++ b/tests/cogs/test_cron.py
@@ -3,6 +3,7 @@ from datetime import datetime
 from unittest.mock import MagicMock, patch
 
 from cogs.cron import Cron
+from cogs.utils.database import _ConnectionContextManager
 
 
 @pytest.fixture
@@ -15,35 +16,31 @@ def mock_bot():
     return bot
 
 
-def create_cog(mock_bot, tmp_path):
+def create_cog(mock_bot):
     with patch.object(Cron, "__init__", lambda self, bot: None):
         with patch("cogs.cron.init_database"):
-            mock_conn = MagicMock()
-            mock_cursor = MagicMock()
-            mock_cursor.fetchall.return_value = []
-            mock_conn.cursor.return_value = mock_cursor
-            mock_conn.close.return_value = None
-
-            with patch("cogs.cron.get_connection", return_value=mock_conn):
-                with patch("cogs.utils.database.get_db_path", return_value=":memory:"):
-                    cog = Cron.__new__(Cron)
-                    cog.bot = mock_bot
-                    cog.scheduled_tasks = []
-                    cog.task = MagicMock()
-                    return cog
+            cog = Cron.__new__(Cron)
+            cog.bot = mock_bot
+            cog.scheduled_tasks = []
+            cog.task = MagicMock()
+            return cog
 
 
 class TestLoadTasks:
-    def test_load_tasks_with_data(self, mock_bot, tmp_path):
+    def test_load_tasks_with_data(self, mock_bot):
         mock_conn = MagicMock()
         mock_cursor = MagicMock()
         mock_cursor.fetchall.return_value = [
             (1, "* * * * *", "audio", 123, 1, None, "test message")
         ]
         mock_conn.cursor.return_value = mock_cursor
-        mock_conn.close.return_value = None
 
-        with patch("cogs.cron.get_connection", return_value=mock_conn):
+        # Create a context manager mock
+        mock_ctx = MagicMock()
+        mock_ctx.__enter__ = MagicMock(return_value=mock_conn)
+        mock_ctx.__exit__ = MagicMock(return_value=False)
+
+        with patch("cogs.cron.get_connection", return_value=mock_ctx):
             cog = Cron.__new__(Cron)
             cog.bot = mock_bot
             cog.scheduled_tasks = []
@@ -52,33 +49,41 @@ class TestLoadTasks:
             assert len(cog.scheduled_tasks) == 1
             assert cog.scheduled_tasks[0]["id"] == 1
 
-
 class TestSaveTasks:
-    def test_save_tasks(self, mock_bot, tmp_path):
-        cog = create_cog(mock_bot, tmp_path)
+    def test_save_tasks(self, mock_bot):
+        cog = create_cog(mock_bot)
         cog.scheduled_tasks = [{"id": 1, "cron_expression": "* * * * *", "type": "audio", "channel_id": 123, "active": True}]
-        cog.save_tasks()
-        assert len(cog.scheduled_tasks) == 1
 
+        mock_conn = MagicMock()
+        mock_cursor = MagicMock()
+        mock_conn.cursor.return_value = mock_cursor
+
+        # Create a context manager mock
+        mock_ctx = MagicMock()
+        mock_ctx.__enter__ = MagicMock(return_value=mock_conn)
+        mock_ctx.__exit__ = MagicMock(return_value=False)
+
+        with patch("cogs.cron.get_connection", return_value=mock_ctx):
+            cog.save_tasks()
+        assert len(cog.scheduled_tasks) == 1
 
 class TestShouldRunTask:
-    def test_should_run_task_invalid_cron(self, mock_bot, tmp_path):
-        cog = create_cog(mock_bot, tmp_path)
+    def test_should_run_task_invalid_cron(self, mock_bot):
+        cog = create_cog(mock_bot)
         task = {"cron_expression": "invalid"}
         now = datetime.now()
         result = cog.should_run_task(task, now)
         assert result is False
 
-    def test_should_run_task_missing_cron_expression(self, mock_bot, tmp_path):
-        cog = create_cog(mock_bot, tmp_path)
+    def test_should_run_task_missing_cron_expression(self, mock_bot):
+        cog = create_cog(mock_bot)
         task = {}
         now = datetime.now()
         result = cog.should_run_task(task, now)
         assert result is False
 
-
 class TestCogUnload:
-    def test_cog_unload_cancels_task(self, mock_bot, tmp_path):
+    def test_cog_unload_cancels_task(self, mock_bot):
         with patch.object(Cron, "__init__", lambda self, bot: None):
             cog = Cron.__new__(Cron)
             cog.bot = mock_bot
@@ -87,10 +92,9 @@ class TestCogUnload:
             cog.cog_unload()
             cog.task.cancel.assert_called_once()
 
-
 class TestTaskManagement:
-    def test_add_task_increments_id(self, mock_bot, tmp_path):
-        cog = create_cog(mock_bot, tmp_path)
+    def test_add_task_increments_id(self, mock_bot):
+        cog = create_cog(mock_bot)
         cog.scheduled_tasks = [
             {"id": 1, "cron_expression": "* * * * *"},
             {"id": 2, "cron_expression": "0 * * * *"},
@@ -99,8 +103,8 @@ class TestTaskManagement:
         new_id = max([t["id"] for t in cog.scheduled_tasks], default=0) + 1
         assert new_id == 3
 
-    def test_find_task_by_id(self, mock_bot, tmp_path):
-        cog = create_cog(mock_bot, tmp_path)
+    def test_find_task_by_id(self, mock_bot):
+        cog = create_cog(mock_bot)
         cog.scheduled_tasks = [
             {"id": 1, "cron_expression": "* * * * *"},
             {"id": 2, "cron_expression": "0 * * * *"},
@@ -110,8 +114,8 @@ class TestTaskManagement:
         assert found_task is not None
         assert found_task["cron_expression"] == "0 * * * *"
 
-    def test_remove_task_by_index(self, mock_bot, tmp_path):
-        cog = create_cog(mock_bot, tmp_path)
+    def test_remove_task_by_index(self, mock_bot):
+        cog = create_cog(mock_bot)
         cog.scheduled_tasks = [
             {"id": 1, "cron_expression": "* * * * *"},
             {"id": 2, "cron_expression": "0 * * * *"},
@@ -121,8 +125,8 @@ class TestTaskManagement:
         assert len(cog.scheduled_tasks) == 1
         assert cog.scheduled_tasks[0]["id"] == 2
 
-    def test_toggle_task_active(self, mock_bot, tmp_path):
-        cog = create_cog(mock_bot, tmp_path)
+    def test_toggle_task_active(self, mock_bot):
+        cog = create_cog(mock_bot)
         cog.scheduled_tasks = [{"id": 1, "active": True}]
 
         cog.scheduled_tasks[0]["active"] = False
@@ -131,14 +135,13 @@ class TestTaskManagement:
         cog.scheduled_tasks[0]["active"] = True
         assert cog.scheduled_tasks[0]["active"] is True
 
-
 class TestTaskTypes:
-    def test_task_type_audio_aliases(self, mock_bot, tmp_path):
+    def test_task_type_audio_aliases(self, mock_bot):
         audio_types = ("zagraj_audio", "audio")
         for task_type in audio_types:
             assert task_type in ("zagraj_audio", "audio")
 
-    def test_task_type_message_aliases(self, mock_bot, tmp_path):
+    def test_task_type_message_aliases(self, mock_bot):
         message_types = ("wyślij_wiadomość", "send_message")
         for task_type in message_types:
-            assert task_type in ("wyślij_wiadomość", "send_message")
\ No newline at end of file
+            assert task_type in ("wyślij_wiadomość", "send_message")