From 572652e712dfc2472cf29b779ef78418becc17d5 Mon Sep 17 00:00:00 2001 From: agessaman Date: Thu, 16 Apr 2026 11:18:59 -0700 Subject: [PATCH] refactor: improve database connection handling in web viewer and tests Updated the BotDataViewer class to utilize a context manager for database connections, enhancing resource management. Additionally, refactored test files to implement a centralized approach for managing SQLite connections, ensuring proper cleanup after tests. This change improves code maintainability and reliability across the application. --- modules/web_viewer/app.py | 20 +++++++------- tests/test_help_command.py | 33 ++++++++++++++++++----- tests/test_maintenance.py | 1 + tests/test_stats_command.py | 26 +++++++++++++++++-- tests/test_web_viewer.py | 52 ++++++++++++++++++++++++++++++++++++- 5 files changed, 113 insertions(+), 19 deletions(-) diff --git a/modules/web_viewer/app.py b/modules/web_viewer/app.py index d9252bd..100020e 100644 --- a/modules/web_viewer/app.py +++ b/modules/web_viewer/app.py @@ -6714,17 +6714,17 @@ class BotDataViewer: return int(time.time() - start_time) else: # Fallback: try to get earliest message timestamp - conn = self._get_db_connection() - cursor = conn.cursor() + with self._with_db_connection() as conn: + cursor = conn.cursor() - # Try to get earliest message timestamp as fallback - cursor.execute(""" - SELECT MIN(timestamp) FROM message_stats - WHERE timestamp IS NOT NULL - """) - result = cursor.fetchone() - if result and result[0]: - return int(time.time() - result[0]) + # Try to get earliest message timestamp as fallback + cursor.execute(""" + SELECT MIN(timestamp) FROM message_stats + WHERE timestamp IS NOT NULL + """) + result = cursor.fetchone() + if result and result[0]: + return int(time.time() - result[0]) return 0 except Exception as e: diff --git a/tests/test_help_command.py b/tests/test_help_command.py index 5d71a76..8051e09 100644 --- a/tests/test_help_command.py +++ b/tests/test_help_command.py @@ -1,9 +1,12 @@ """Tests for modules.commands.help_command — pure logic and integration paths.""" import configparser +import sqlite3 from contextlib import contextmanager from unittest.mock import MagicMock, Mock +import pytest + from modules.commands.help_command import HelpCommand from tests.conftest import mock_message @@ -11,6 +14,27 @@ from tests.conftest import mock_message # Bot factory # --------------------------------------------------------------------------- +_TRACKED_CONNECTIONS = [] + + +def _create_tracked_connection(): + conn = sqlite3.connect(":memory:") + _TRACKED_CONNECTIONS.append(conn) + return conn + + +@pytest.fixture(autouse=True) +def _close_tracked_connections(): + """Ensure every test-created sqlite connection is closed.""" + yield + while _TRACKED_CONNECTIONS: + conn = _TRACKED_CONNECTIONS.pop() + try: + conn.close() + except sqlite3.Error: + pass + + def _make_bot(enabled=True, commands=None): """Create a minimal mock bot for HelpCommand tests.""" bot = MagicMock() @@ -46,8 +70,7 @@ def _make_bot(enabled=True, commands=None): bot.command_manager.plugin_loader.keyword_mappings = {} # DB manager with in-memory SQLite - import sqlite3 - conn = sqlite3.connect(":memory:") + conn = _create_tracked_connection() db = MagicMock() @contextmanager @@ -310,11 +333,10 @@ class TestGetAvailableCommandsListFiltered: def test_command_in_stats_not_in_keyword_mappings(self): """Commands returned from DB stats but not in keyword_mappings (lines 218-235).""" - import sqlite3 from contextlib import contextmanager bot = _make_bot() - conn = sqlite3.connect(":memory:") + conn = _create_tracked_connection() conn.execute(""" CREATE TABLE command_stats ( id INTEGER PRIMARY KEY, @@ -422,11 +444,10 @@ class TestGetAvailableCommandsList: def test_with_stats_table_present(self): """When command_stats table exists, commands are sorted by usage count.""" - import sqlite3 from contextlib import contextmanager bot = _make_bot() - conn = sqlite3.connect(":memory:") + conn = _create_tracked_connection() conn.execute(""" CREATE TABLE command_stats ( id INTEGER PRIMARY KEY, diff --git a/tests/test_maintenance.py b/tests/test_maintenance.py index 3189a2a..0d75002 100644 --- a/tests/test_maintenance.py +++ b/tests/test_maintenance.py @@ -75,6 +75,7 @@ class TestRowN: cur = _Cur() assert _row_n(cur) == 7 + conn.close() # --------------------------------------------------------------------------- diff --git a/tests/test_stats_command.py b/tests/test_stats_command.py index d0997c8..9c68b83 100644 --- a/tests/test_stats_command.py +++ b/tests/test_stats_command.py @@ -1,17 +1,39 @@ """Tests for modules.commands.stats_command — pure logic functions.""" import configparser +import sqlite3 from contextlib import contextmanager from unittest.mock import MagicMock, Mock, patch +import pytest + from modules.commands.stats_command import StatsCommand from tests.conftest import mock_message +_TRACKED_CONNECTIONS = [] + + +def _create_tracked_connection(): + conn = sqlite3.connect(":memory:") + _TRACKED_CONNECTIONS.append(conn) + return conn + + +@pytest.fixture(autouse=True) +def _close_tracked_connections(): + """Ensure each test closes its sqlite connections.""" + yield + while _TRACKED_CONNECTIONS: + conn = _TRACKED_CONNECTIONS.pop() + try: + conn.close() + except sqlite3.Error: + pass + def _make_db_manager(): """Create a mock db_manager with a working connection context manager.""" - import sqlite3 - conn = sqlite3.connect(":memory:") + conn = _create_tracked_connection() db = MagicMock() @contextmanager diff --git a/tests/test_web_viewer.py b/tests/test_web_viewer.py index 8f0458f..2a87d8f 100644 --- a/tests/test_web_viewer.py +++ b/tests/test_web_viewer.py @@ -116,6 +116,26 @@ def auth_client(auth_viewer): yield c +@pytest.fixture(autouse=True) +def cleanup_sqlite_connections(monkeypatch): + """Track and close SQLite connections opened during each test.""" + tracked_connections = [] + original_connect = sqlite3.connect + + def _tracked_connect(*args, **kwargs): + conn = original_connect(*args, **kwargs) + tracked_connections.append(conn) + return conn + + monkeypatch.setattr(sqlite3, "connect", _tracked_connect) + yield + for conn in tracked_connections: + try: + conn.close() + except sqlite3.Error: + pass + + # --------------------------------------------------------------------------- # Helper: insert a contact row so contact-related routes have data # --------------------------------------------------------------------------- @@ -1660,9 +1680,37 @@ def socketio_viewer(tmp_path_factory): v.app.config["TESTING"] = True v.app.config["SECRET_KEY"] = "test-secret" - return v + yield v + with v._clients_lock: + v.connected_clients.clear() +@pytest.fixture +def managed_socketio_clients(monkeypatch, socketio_viewer): + """Track SocketIO test clients and always disconnect on teardown.""" + import flask_socketio + + created_clients = [] + original_client_cls = flask_socketio.SocketIOTestClient + + class ManagedSocketIOTestClient(original_client_cls): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + created_clients.append(self) + + monkeypatch.setattr(flask_socketio, "SocketIOTestClient", ManagedSocketIOTestClient) + yield + for client in created_clients: + try: + if client.is_connected(): + client.disconnect() + except Exception: + pass + with socketio_viewer._clients_lock: + assert not socketio_viewer.connected_clients + + +@pytest.mark.usefixtures("managed_socketio_clients") class TestSubscribeCommandsHistoryReplay: """subscribe_commands must replay last 50 command rows on connect (TASK-02 / BUG-023).""" @@ -2615,6 +2663,7 @@ class TestMaintenanceStatusFields: # T1-A: subscribe_packets and subscribe_messages history replay # =========================================================================== +@pytest.mark.usefixtures("managed_socketio_clients") class TestSubscribePacketsHistoryReplay: """subscribe_packets must replay last 50 packet/command/routing rows on connect (T1-A).""" @@ -2763,6 +2812,7 @@ class TestSubscribePacketsHistoryReplay: assert seq_values == [0, 1, 2] +@pytest.mark.usefixtures("managed_socketio_clients") class TestSubscribeMessagesHistoryReplay: """subscribe_messages must replay last 50 message rows on connect (T1-A)."""