From 3978973ed1e1e904275509086972c3bf1800ced4 Mon Sep 17 00:00:00 2001 From: Orinks <38449772+Orinks@users.noreply.github.com> Date: Mon, 2 Mar 2026 16:22:54 +0000 Subject: [PATCH] fix(sftp): wire asyncssh progress_handler for transfer progress (#45) After the asyncssh migration, SFTP downloads/uploads showed 0% progress because the manual chunk-read loop bypassed asyncssh's pipelined I/O and native progress reporting. Switch SFTPClient.download() and upload() to use sftp.get()/sftp.put() with their progress_handler callback. This gives pipelined reads/writes for better throughput and correct progress updates via the existing wx.PostEvent thread-safe notification path. A chunked-read fallback is retained for in-memory BinaryIO streams (BytesIO) used in tests. Co-Authored-By: Claude Opus 4.6 --- src/portkeydrop/protocols.py | 99 ++++++++++++------ tests/test_protocols.py | 197 +++++++++++++++++++++++++++++++++++ 2 files changed, 265 insertions(+), 31 deletions(-) diff --git a/src/portkeydrop/protocols.py b/src/portkeydrop/protocols.py index 4efa4fc..e0b3215 100644 --- a/src/portkeydrop/protocols.py +++ b/src/portkeydrop/protocols.py @@ -657,43 +657,80 @@ def download( self, remote_path: str, local_file: BinaryIO, callback: ProgressCallback | None = None ) -> None: sftp = self._ensure_connected() - - async def _download(): - async with sftp.open(remote_path, "rb") as rf: - total = (await sftp.stat(remote_path)).size or 0 - transferred = 0 - while True: - chunk = await rf.read(8192) - if not chunk: - break - local_file.write(chunk) - transferred += len(chunk) - if callback: - callback(transferred, total) - - self._run(_download()) + local_path = getattr(local_file, "name", None) + + if isinstance(local_path, str) and os.path.isabs(local_path): + # asyncssh native get() — pipelined reads with progress reporting + local_file.close() + + async def _download(): + handler = None + if callback: + + def handler(srcpath, dstpath, copied, total): + callback(copied, total) + + await sftp.get(remote_path, local_path, progress_handler=handler) + + self._run(_download()) + else: + # Fallback for in-memory streams (BytesIO, etc.) + async def _download(): + async with sftp.open(remote_path, "rb") as rf: + total = (await sftp.stat(remote_path)).size or 0 + transferred = 0 + while True: + chunk = await rf.read(8192) + if not chunk: + break + local_file.write(chunk) + transferred += len(chunk) + if callback: + callback(transferred, total) + + self._run(_download()) def upload( self, local_file: BinaryIO, remote_path: str, callback: ProgressCallback | None = None ) -> None: sftp = self._ensure_connected() - local_file.seek(0, 2) - total = local_file.tell() - local_file.seek(0) + local_path = getattr(local_file, "name", None) + + if isinstance(local_path, str) and os.path.isabs(local_path): + # asyncssh native put() — pipelined writes with progress reporting + total = os.path.getsize(local_path) + local_file.close() + + async def _upload(): + handler = None + if callback: + + def handler(srcpath, dstpath, copied, total_bytes): + callback(copied, total_bytes) + + await sftp.put(local_path, remote_path, progress_handler=handler) + + self._run(_upload()) + else: + # Fallback for in-memory streams (BytesIO, etc.) + local_file.seek(0, 2) + total = local_file.tell() + local_file.seek(0) + + async def _upload(): + async with sftp.open(remote_path, "wb") as wf: + transferred = 0 + while True: + chunk = local_file.read(8192) + if not chunk: + break + await wf.write(chunk) + transferred += len(chunk) + if callback: + callback(transferred, total) + + self._run(_upload()) - async def _upload(): - async with sftp.open(remote_path, "wb") as wf: - transferred = 0 - while True: - chunk = local_file.read(8192) - if not chunk: - break - await wf.write(chunk) - transferred += len(chunk) - if callback: - callback(transferred, total) - - self._run(_upload()) remote_attrs = self._run(sftp.stat(remote_path)) remote_size = remote_attrs.size or 0 if remote_size != total: diff --git a/tests/test_protocols.py b/tests/test_protocols.py index 7be6d81..72c1c88 100644 --- a/tests/test_protocols.py +++ b/tests/test_protocols.py @@ -668,6 +668,203 @@ def test_parent_from_root(self): assert result == "/" +class TestSFTPClientNativeTransfer: + """Tests for asyncssh sftp.get()/put() with progress_handler.""" + + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_download_uses_native_get_with_progress(self, mock_connect): + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.return_value = "/" + mock_conn.start_sftp_client.return_value = mock_sftp + mock_connect.return_value = mock_conn + + progress_calls: list[tuple[int, int]] = [] + + async def fake_get(remotepath, localpath, *, progress_handler=None, **kwargs): + if progress_handler: + progress_handler(remotepath, localpath, 0, 1000) + progress_handler(remotepath, localpath, 500, 1000) + progress_handler(remotepath, localpath, 1000, 1000) + + mock_sftp.get = AsyncMock(side_effect=fake_get) + + client = SFTPClient(ConnectionInfo(protocol=Protocol.SFTP, host="example.com")) + client.connect() + + close_mock = MagicMock() + mock_file = MagicMock() + mock_file.name = "/tmp/downloaded.bin" + mock_file.close = close_mock + + client.download( + "/remote/file.bin", + mock_file, + callback=lambda t, n: progress_calls.append((t, n)), + ) + + mock_sftp.get.assert_awaited_once() + call_kwargs = mock_sftp.get.call_args + assert call_kwargs[0][0] == "/remote/file.bin" + assert call_kwargs[0][1] == "/tmp/downloaded.bin" + assert call_kwargs[1]["progress_handler"] is not None + assert progress_calls == [(0, 1000), (500, 1000), (1000, 1000)] + close_mock.assert_called_once() + + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_download_no_callback_still_uses_native_get(self, mock_connect): + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.return_value = "/" + mock_conn.start_sftp_client.return_value = mock_sftp + mock_connect.return_value = mock_conn + + client = SFTPClient(ConnectionInfo(protocol=Protocol.SFTP, host="example.com")) + client.connect() + + close_mock = MagicMock() + mock_file = MagicMock() + mock_file.name = "/tmp/downloaded.bin" + mock_file.close = close_mock + + client.download("/remote/file.bin", mock_file) + + mock_sftp.get.assert_awaited_once() + call_kwargs = mock_sftp.get.call_args + assert call_kwargs[1]["progress_handler"] is None + close_mock.assert_called_once() + + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_download_bytesio_uses_chunked_fallback(self, mock_connect): + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.return_value = "/" + mock_conn.start_sftp_client.return_value = mock_sftp + mock_connect.return_value = mock_conn + + mock_remote_file = AsyncMock() + mock_remote_file.read.side_effect = [b"hello", b""] + mock_open_cm = MagicMock() + mock_open_cm.__aenter__ = AsyncMock(return_value=mock_remote_file) + mock_open_cm.__aexit__ = AsyncMock(return_value=False) + mock_sftp.open = MagicMock(return_value=mock_open_cm) + stat_attrs = MagicMock() + stat_attrs.size = 5 + mock_sftp.stat.return_value = stat_attrs + + client = SFTPClient(ConnectionInfo(protocol=Protocol.SFTP, host="example.com")) + client.connect() + + buf = io.BytesIO() + progress_calls: list[tuple[int, int]] = [] + client.download( + "/remote.txt", + buf, + callback=lambda t, n: progress_calls.append((t, n)), + ) + + assert buf.getvalue() == b"hello" + assert progress_calls == [(5, 5)] + mock_sftp.get.assert_not_awaited() + + @patch("os.path.getsize", return_value=4) + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_upload_uses_native_put_with_progress(self, mock_connect, _mock_getsize): + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.return_value = "/" + mock_conn.start_sftp_client.return_value = mock_sftp + mock_connect.return_value = mock_conn + + progress_calls: list[tuple[int, int]] = [] + + async def fake_put(localpath, remotepath, *, progress_handler=None, **kwargs): + if progress_handler: + progress_handler(localpath, remotepath, 0, 4) + progress_handler(localpath, remotepath, 4, 4) + + mock_sftp.put = AsyncMock(side_effect=fake_put) + + stat_attrs = MagicMock() + stat_attrs.size = 4 + mock_sftp.stat.return_value = stat_attrs + + client = SFTPClient(ConnectionInfo(protocol=Protocol.SFTP, host="example.com")) + client.connect() + + close_mock = MagicMock() + mock_file = MagicMock() + mock_file.name = "/tmp/upload.bin" + mock_file.close = close_mock + + client.upload( + mock_file, + "/remote/file.bin", + callback=lambda t, n: progress_calls.append((t, n)), + ) + + mock_sftp.put.assert_awaited_once() + call_kwargs = mock_sftp.put.call_args + assert call_kwargs[0][0] == "/tmp/upload.bin" + assert call_kwargs[0][1] == "/remote/file.bin" + assert call_kwargs[1]["progress_handler"] is not None + assert progress_calls == [(0, 4), (4, 4)] + close_mock.assert_called_once() + + @patch("os.path.getsize", return_value=4) + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_upload_verifies_remote_size(self, mock_connect, _mock_getsize): + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.return_value = "/" + mock_conn.start_sftp_client.return_value = mock_sftp + mock_connect.return_value = mock_conn + + stat_attrs = MagicMock() + stat_attrs.size = 2 # Mismatch: expected 4 + mock_sftp.stat.return_value = stat_attrs + + client = SFTPClient(ConnectionInfo(protocol=Protocol.SFTP, host="example.com")) + client.connect() + + mock_file = MagicMock() + mock_file.name = "/tmp/upload.bin" + + with pytest.raises(RuntimeError, match="verification failed"): + client.upload(mock_file, "/remote.bin") + + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_upload_bytesio_uses_chunked_fallback(self, mock_connect): + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.return_value = "/" + mock_conn.start_sftp_client.return_value = mock_sftp + mock_connect.return_value = mock_conn + + mock_write_file = AsyncMock() + mock_open_cm = MagicMock() + mock_open_cm.__aenter__ = AsyncMock(return_value=mock_write_file) + mock_open_cm.__aexit__ = AsyncMock(return_value=False) + mock_sftp.open = MagicMock(return_value=mock_open_cm) + + stat_attrs = MagicMock() + stat_attrs.size = 4 + mock_sftp.stat.return_value = stat_attrs + + client = SFTPClient(ConnectionInfo(protocol=Protocol.SFTP, host="example.com")) + client.connect() + + progress_calls: list[tuple[int, int]] = [] + client.upload( + io.BytesIO(b"data"), + "/remote.txt", + callback=lambda t, n: progress_calls.append((t, n)), + ) + + assert progress_calls == [(4, 4)] + mock_sftp.put.assert_not_awaited() + + class TestProtocolEnum: def test_all_protocols(self): assert Protocol.FTP.value == "ftp"