diff --git a/src/portkeydrop/dialogs/transfer.py b/src/portkeydrop/dialogs/transfer.py index 82bcc44..fb7e15c 100644 --- a/src/portkeydrop/dialogs/transfer.py +++ b/src/portkeydrop/dialogs/transfer.py @@ -165,7 +165,8 @@ def callback(transferred: int, total: int) -> None: if item.cancel_event.is_set(): raise InterruptedError("Transfer cancelled") item.transferred_bytes = transferred - item.total_bytes = total + if total > 0: + item.total_bytes = total self._notify() client.download(item.remote_path, f, callback=callback) @@ -187,7 +188,8 @@ def callback(transferred: int, total: int) -> None: if item.cancel_event.is_set(): raise InterruptedError("Transfer cancelled") item.transferred_bytes = transferred - item.total_bytes = total + if total > 0: + item.total_bytes = total self._notify() client.upload(f, item.remote_path, callback=callback) @@ -208,6 +210,15 @@ def _run_recursive_download(self, client: TransferClient, item: TransferItem) -> # Collect all files first to calculate total size file_queue: list[tuple[str, str, int]] = [] # (remote, local, size) self._collect_remote_files(client, item.remote_path, item.local_path, file_queue) + # Re-stat files with size=0 to resolve symlink targets + for i, (remote_file, local_file, size) in enumerate(file_queue): + if size == 0: + try: + real_size = client.stat(remote_file).size + if real_size > 0: + file_queue[i] = (remote_file, local_file, real_size) + except Exception: + pass item.total_bytes = sum(size for _, _, size in file_queue) item.transferred_bytes = 0 self._notify() diff --git a/src/portkeydrop/protocols.py b/src/portkeydrop/protocols.py index 2eefe1d..e7a9b35 100644 --- a/src/portkeydrop/protocols.py +++ b/src/portkeydrop/protocols.py @@ -768,6 +768,12 @@ def download( sftp = self._ensure_connected() local_path = getattr(local_file, "name", None) + # Resolve symlinks so stat() returns the real file size + try: + resolved = self._run(sftp.realpath(remote_path)) + except Exception: + resolved = remote_path + if isinstance(local_path, str) and os.path.isabs(local_path): # asyncssh native get() — pipelined reads with progress reporting local_file.close() @@ -779,14 +785,14 @@ async def _download(): def handler(srcpath, dstpath, copied, total): callback(copied, total) - await sftp.get(remote_path, local_path, progress_handler=handler) + await sftp.get(resolved, 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 + async with sftp.open(resolved, "rb") as rf: + total = (await sftp.stat(resolved)).size or 0 transferred = 0 while True: chunk = await rf.read(8192) diff --git a/tests/test_protocols.py b/tests/test_protocols.py index e050b04..072f3f4 100644 --- a/tests/test_protocols.py +++ b/tests/test_protocols.py @@ -452,7 +452,7 @@ def test_list_dir_maps_file_attributes(self, mock_connect): def test_chdir_download_upload_and_file_ops(self, mock_connect): mock_conn = AsyncMock() mock_sftp = AsyncMock() - mock_sftp.realpath.side_effect = ["/", "/uploads"] + mock_sftp.realpath.side_effect = ["/", "/uploads", "/remote.bin"] # chdir now validates with stat — return directory attributes chdir_stat_attrs = MagicMock() chdir_stat_attrs.permissions = stat_mod.S_IFDIR | 0o755 @@ -688,7 +688,7 @@ class TestSFTPClientNativeTransfer: def test_download_uses_native_get_with_progress(self, mock_connect): mock_conn = AsyncMock() mock_sftp = AsyncMock() - mock_sftp.realpath.return_value = "/" + mock_sftp.realpath.side_effect = lambda p: "/" if p == "." else p mock_conn.start_sftp_client.return_value = mock_sftp mock_connect.return_value = mock_conn @@ -728,7 +728,7 @@ async def fake_get(remotepath, localpath, *, progress_handler=None, **kwargs): def test_download_no_callback_still_uses_native_get(self, mock_connect): mock_conn = AsyncMock() mock_sftp = AsyncMock() - mock_sftp.realpath.return_value = "/" + mock_sftp.realpath.side_effect = lambda p: "/" if p == "." else p mock_conn.start_sftp_client.return_value = mock_sftp mock_connect.return_value = mock_conn @@ -751,7 +751,7 @@ def test_download_no_callback_still_uses_native_get(self, mock_connect): def test_download_bytesio_uses_chunked_fallback(self, mock_connect): mock_conn = AsyncMock() mock_sftp = AsyncMock() - mock_sftp.realpath.return_value = "/" + mock_sftp.realpath.side_effect = lambda p: "/" if p == "." else p mock_conn.start_sftp_client.return_value = mock_sftp mock_connect.return_value = mock_conn @@ -878,6 +878,126 @@ def test_upload_bytesio_uses_chunked_fallback(self, mock_connect): mock_sftp.put.assert_not_awaited() +class TestSFTPDownloadSymlinkResolution: + """Tests that SFTPClient.download() resolves symlinks via realpath.""" + + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_download_resolves_symlink_via_realpath(self, mock_connect): + """Native get() path uses the resolved path for symlinked files.""" + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.side_effect = [ + "/", # connect + "/real/file.bin", # download resolves symlink + ] + 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, 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() + + mock_file = MagicMock() + mock_file.name = "/tmp/downloaded.bin" + mock_file.close = MagicMock() + + client.download( + "/symlink/file.bin", + mock_file, + callback=lambda t, n: progress_calls.append((t, n)), + ) + + # Verify get() was called with the resolved path + call_args = mock_sftp.get.call_args + assert call_args[0][0] == "/real/file.bin" + assert progress_calls == [(500, 1000), (1000, 1000)] + + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_download_fallback_resolves_symlink_via_realpath(self, mock_connect): + """BytesIO fallback path uses the resolved path for symlinked files.""" + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.side_effect = [ + "/", # connect + "/real/file.txt", # download resolves symlink + ] + 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( + "/symlink/file.txt", + buf, + callback=lambda t, n: progress_calls.append((t, n)), + ) + + # Verify open() and stat() were called with the resolved path + mock_sftp.open.assert_called_once_with("/real/file.txt", "rb") + mock_sftp.stat.assert_called_once_with("/real/file.txt") + assert buf.getvalue() == b"hello" + assert progress_calls == [(5, 5)] + + @patch("asyncssh.connect", new_callable=AsyncMock) + def test_download_falls_back_to_original_path_when_realpath_fails(self, mock_connect): + """If realpath fails, download uses the original path.""" + mock_conn = AsyncMock() + mock_sftp = AsyncMock() + mock_sftp.realpath.side_effect = [ + "/", # connect + OSError("realpath failed"), # download fallback + ] + 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, 100, 100) + + mock_sftp.get = AsyncMock(side_effect=fake_get) + + client = SFTPClient(ConnectionInfo(protocol=Protocol.SFTP, host="example.com")) + client.connect() + + mock_file = MagicMock() + mock_file.name = "/tmp/downloaded.bin" + mock_file.close = MagicMock() + + client.download( + "/original/path.bin", + mock_file, + callback=lambda t, n: progress_calls.append((t, n)), + ) + + # Verify get() was called with the original path (fallback) + call_args = mock_sftp.get.call_args + assert call_args[0][0] == "/original/path.bin" + assert progress_calls == [(100, 100)] + + class TestProtocolEnum: def test_all_protocols(self): assert Protocol.FTP.value == "ftp" diff --git a/tests/test_transfer_symlink.py b/tests/test_transfer_symlink.py new file mode 100644 index 0000000..315d713 --- /dev/null +++ b/tests/test_transfer_symlink.py @@ -0,0 +1,200 @@ +"""Tests for symlink-related download progress fixes (#52).""" + +from __future__ import annotations + +import io +from unittest.mock import MagicMock, patch + + +from portkeydrop.dialogs.transfer import TransferManager, TransferStatus +from portkeydrop.protocols import RemoteFile + + +class TestCallbackTotalBytesGuard: + """Fix 2: callback must not overwrite item.total_bytes with 0.""" + + def test_download_callback_preserves_total_bytes_when_zero_reported(self): + """If the progress callback reports total=0, item.total_bytes should + not be overwritten when it already holds a positive value.""" + mock_client = MagicMock() + captured_callback = {} + + def fake_download(remote_path, local_file, callback=None): + captured_callback["fn"] = callback + if callback: + # First call reports real total + callback(500, 1000) + # Subsequent call reports total=0 (symlink stat issue) + callback(600, 0) + + mock_client.download.side_effect = fake_download + + manager = TransferManager(notify_window=None) + + with patch("builtins.open", return_value=MagicMock(spec=io.BufferedWriter)): + item = manager.add_download(mock_client, "/remote/file.bin", "/tmp/file.bin", 1000) + # Wait for the thread to finish + import time + + deadline = time.monotonic() + 5 + while item.status == TransferStatus.IN_PROGRESS and time.monotonic() < deadline: + time.sleep(0.05) + + assert item.total_bytes == 1000 + assert item.transferred_bytes == 600 + assert item.status == TransferStatus.COMPLETED + + def test_upload_callback_preserves_total_bytes_when_zero_reported(self): + """Upload callback should also guard total_bytes from being zeroed.""" + mock_client = MagicMock() + + def fake_upload(local_file, remote_path, callback=None): + if callback: + callback(200, 500) + callback(400, 0) # symlink-related 0 + + mock_client.upload.side_effect = fake_upload + + manager = TransferManager(notify_window=None) + + with patch("builtins.open", return_value=MagicMock(spec=io.BufferedReader)): + item = manager.add_upload(mock_client, "/tmp/file.bin", "/remote/file.bin", 500) + import time + + deadline = time.monotonic() + 5 + while item.status == TransferStatus.IN_PROGRESS and time.monotonic() < deadline: + time.sleep(0.05) + + assert item.total_bytes == 500 + assert item.transferred_bytes == 400 + assert item.status == TransferStatus.COMPLETED + + def test_download_callback_updates_total_bytes_when_positive(self): + """Callback should update total_bytes when a positive value is reported.""" + mock_client = MagicMock() + + def fake_download(remote_path, local_file, callback=None): + if callback: + callback(100, 2000) + + mock_client.download.side_effect = fake_download + + manager = TransferManager(notify_window=None) + + with patch("builtins.open", return_value=MagicMock(spec=io.BufferedWriter)): + item = manager.add_download(mock_client, "/remote/file.bin", "/tmp/file.bin", 0) + import time + + deadline = time.monotonic() + 5 + while item.status == TransferStatus.IN_PROGRESS and time.monotonic() < deadline: + time.sleep(0.05) + + assert item.total_bytes == 2000 + assert item.status == TransferStatus.COMPLETED + + +class TestRecursiveDownloadRestat: + """Fix 3: recursive download re-stats files with size=0.""" + + def test_recursive_download_restats_zero_size_files(self): + """Files with size=0 in the listing should be individually re-statted + to resolve symlink targets and get the real size.""" + mock_client = MagicMock() + + # list_dir returns one normal file and one symlinked file (size=0) + mock_client.list_dir.return_value = [ + RemoteFile(name="normal.txt", path="/remote/dir/normal.txt", size=500), + RemoteFile(name="symlink.txt", path="/remote/dir/symlink.txt", size=0), + ] + + # stat on the symlinked file returns the real size + mock_client.stat.return_value = RemoteFile( + name="symlink.txt", path="/remote/dir/symlink.txt", size=750 + ) + + def fake_download(remote_path, local_file, callback=None): + if callback: + size = 500 if "normal" in remote_path else 750 + callback(size, size) + + mock_client.download.side_effect = fake_download + + manager = TransferManager(notify_window=None) + + with patch("builtins.open", return_value=MagicMock(spec=io.BufferedWriter)): + with patch("os.makedirs"): + item = manager.add_recursive_download(mock_client, "/remote/dir", "/tmp/local_dir") + import time + + deadline = time.monotonic() + 5 + while item.status == TransferStatus.IN_PROGRESS and time.monotonic() < deadline: + time.sleep(0.05) + + # Total should be 500 + 750 = 1250 (not 500 + 0 = 500) + assert item.total_bytes == 1250 + assert item.status == TransferStatus.COMPLETED + # stat should have been called for the zero-size file + mock_client.stat.assert_called_once_with("/remote/dir/symlink.txt") + + def test_recursive_download_skips_restat_for_nonzero_files(self): + """Files with size > 0 should not be re-statted.""" + mock_client = MagicMock() + + mock_client.list_dir.return_value = [ + RemoteFile(name="file1.txt", path="/remote/dir/file1.txt", size=100), + RemoteFile(name="file2.txt", path="/remote/dir/file2.txt", size=200), + ] + + def fake_download(remote_path, local_file, callback=None): + if callback: + size = 100 if "file1" in remote_path else 200 + callback(size, size) + + mock_client.download.side_effect = fake_download + + manager = TransferManager(notify_window=None) + + with patch("builtins.open", return_value=MagicMock(spec=io.BufferedWriter)): + with patch("os.makedirs"): + item = manager.add_recursive_download(mock_client, "/remote/dir", "/tmp/local_dir") + import time + + deadline = time.monotonic() + 5 + while item.status == TransferStatus.IN_PROGRESS and time.monotonic() < deadline: + time.sleep(0.05) + + assert item.total_bytes == 300 + assert item.status == TransferStatus.COMPLETED + # stat should NOT have been called (no zero-size files) + mock_client.stat.assert_not_called() + + def test_recursive_download_handles_stat_failure_gracefully(self): + """If stat fails on a zero-size file, it should remain at size=0.""" + mock_client = MagicMock() + + mock_client.list_dir.return_value = [ + RemoteFile(name="broken_link.txt", path="/remote/dir/broken_link.txt", size=0), + ] + + mock_client.stat.side_effect = OSError("No such file") + + def fake_download(remote_path, local_file, callback=None): + if callback: + callback(0, 0) + + mock_client.download.side_effect = fake_download + + manager = TransferManager(notify_window=None) + + with patch("builtins.open", return_value=MagicMock(spec=io.BufferedWriter)): + with patch("os.makedirs"): + item = manager.add_recursive_download(mock_client, "/remote/dir", "/tmp/local_dir") + import time + + deadline = time.monotonic() + 5 + while item.status == TransferStatus.IN_PROGRESS and time.monotonic() < deadline: + time.sleep(0.05) + + assert item.total_bytes == 0 + assert item.status == TransferStatus.COMPLETED + mock_client.stat.assert_called_once_with("/remote/dir/broken_link.txt")