From ee81ef18996c0540e2dd3d3a50df3689222e5f58 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Mon, 10 Nov 2025 01:37:53 +0000 Subject: [PATCH] Fix: Correct asyncssh close pattern and test suite errors - Patched `SSHManager.disconnect` to use the correct `conn.close()` and `await conn.wait_closed()` pattern, resolving a `TypeError`. - Added `pytest-mock` to the project dependencies to fix the "mocker not found" test error. - Refactored `test_run_command_always_closes_connection` to use a more accurate mock for `asyncssh` streams, resolving an `AttributeError`. - Fixed `test_apply_update_failure_and_rollback` by creating dummy files and mocking `shutil.copy`, ensuring the test correctly validates the rollback logic. --- pyproject.toml | 1 + src/ssh_manager.py | 24 ++++++++++++++--------- tests/test_ssh_manager.py | 40 +++++++++++++++++++++++++++------------ tests/test_updater.py | 29 +++++++++++++++++++++------- 4 files changed, 66 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8e7cdf9..08cf744 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ "asyncssh", "pytest", "pytest-asyncio", + "pytest-mock", "tenacity", "async-timeout", "requests", diff --git a/src/ssh_manager.py b/src/ssh_manager.py index a5d71db..fbea8ab 100644 --- a/src/ssh_manager.py +++ b/src/ssh_manager.py @@ -177,16 +177,22 @@ async def run_command_in_shell(self, alias: str, command: str) -> str: async def disconnect(self, alias: str): """ - Closes a persistent shell connection. + Safely closes a persistent shell connection. """ - if alias in self.active_shells: - logger.info(f"Closing interactive shell for {alias}.") - conn = self.active_shells[alias] - with contextlib.suppress(Exception): - conn.close() - if hasattr(conn, "wait_closed"): - await conn.wait_closed() - del self.active_shells[alias] + conn = self.active_shells.get(alias) + if not conn: + return # No active connection to close + + logger.info(f"Closing interactive shell for {alias}...") + try: + # AsyncSSH close() is synchronous, but wait_closed() is awaitable. + conn.close() + if hasattr(conn, "wait_closed"): + await conn.wait_closed() + except Exception as e: + logger.warning(f"Error while closing SSH session for {alias}: {e}", exc_info=True) + finally: + self.active_shells.pop(alias, None) async def close_all_connections(self): """Closes all active shell connections.""" diff --git a/tests/test_ssh_manager.py b/tests/test_ssh_manager.py index b8ad0fa..bfc9a0d 100644 --- a/tests/test_ssh_manager.py +++ b/tests/test_ssh_manager.py @@ -98,14 +98,29 @@ async def test_run_command_always_closes_connection(mocker): mocker.patch.object(manager, '_create_connection', return_value=mock_conn) mocker.patch.object(manager, '_close_conn', new_callable=AsyncMock) - # Mock the command streaming part - async def fake_streamer(*args, **kwargs): - yield ("output", "stdout") - - # Mock the asyncssh process and its streams + # Mock the asyncssh process and its streams to behave like real streams process_mock = AsyncMock() - process_mock.stdout = fake_streamer() - process_mock.stderr = fake_streamer() + + # Mock stdout to be an async iterator and have a readline method + stdout_mock = AsyncMock() + stdout_mock.readline = AsyncMock(return_value="12345") # Mock PID + + async def fake_stdout_stream(): + yield "output line 1" + + # Make the mock iterable for the `async for` loop + stdout_mock.__aiter__ = MagicMock(return_value=fake_stdout_stream()) + + # Mock stderr to be an async iterator + stderr_mock = AsyncMock() + + async def fake_stderr_stream(): + yield "error line 1" + + stderr_mock.__aiter__ = MagicMock(return_value=fake_stderr_stream()) + + process_mock.stdout = stdout_mock + process_mock.stderr = stderr_mock conn_mock = AsyncMock() conn_mock.create_process.return_value = process_mock @@ -123,14 +138,15 @@ async def fake_streamer(*args, **kwargs): # 2. Test exception case manager._close_conn.reset_mock() - # Configure streamer to raise an error - async def error_streamer(*args, **kwargs): - yield ("output", "stdout") + # Configure the mocked stdout stream to raise an error during iteration + async def error_stream(): + yield "output line 1" # The stream must yield something before failing raise ValueError("Command failed") - process_mock.stdout = error_streamer() + # Replace the iterator of the existing mock, preserving the `readline` method + stdout_mock.__aiter__ = MagicMock(return_value=error_stream()) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Command failed"): async for _, __ in manager.run_command("alias", "cmd"): pass diff --git a/tests/test_updater.py b/tests/test_updater.py index 1ddf740..81edae0 100644 --- a/tests/test_updater.py +++ b/tests/test_updater.py @@ -29,19 +29,34 @@ def test_apply_update_success(mock_sleep, mock_run, mock_download, mock_copytree assert mock_run.call_count == 3 @patch('src.updater.shutil.rmtree') +@patch('src.updater.shutil.copy') @patch('src.updater.shutil.copytree') @patch('src.updater.download_and_extract_zip', side_effect=Exception("Download failed")) @patch('src.updater.rollback') @patch('src.updater.time.sleep', return_value=None) -def test_apply_update_failure_and_rollback(mock_sleep, mock_rollback, mock_download, mock_copytree, mock_rmtree): +def test_apply_update_failure_and_rollback(mock_sleep, mock_rollback, mock_download, mock_copytree, mock_copy, mock_rmtree): """Test a failed update and the subsequent rollback.""" - with patch('src.updater.subprocess.run') as mock_run: - result = apply_update() + # Create dummy data files for the backup step to succeed before the download fails + from pathlib import Path + config_file = Path("config.json") + db_file = Path("database.db") + config_file.touch() + db_file.touch() + + try: + with patch('src.updater.subprocess.run') as mock_run: + result = apply_update() - mock_run.assert_called_once_with(["systemctl", "stop", "telegram_bot.service"], check=True, timeout=COMMAND_TIMEOUT) - assert "Update Failed: Download failed" in result - assert "Attempting to roll back..." in result - mock_rollback.assert_called_once() + mock_run.assert_called_once_with(["systemctl", "stop", "telegram_bot.service"], check=True, timeout=COMMAND_TIMEOUT) + assert "Update Failed: Download failed" in result + assert "Attempting to roll back..." in result + mock_rollback.assert_called_once() + finally: + # Clean up the dummy files to not affect other tests + if config_file.exists(): + config_file.unlink() + if db_file.exists(): + db_file.unlink() @patch('src.updater.shutil.copytree') @patch('src.updater.subprocess.run')