Fix: Correct asyncssh close pattern and test suite errors

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "asyncssh",
     "pytest",
     "pytest-asyncio",
+    "pytest-mock",
     "tenacity",
     "async-timeout",
     "requests",

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."""

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
 

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')