diff --git a/aexpect/remote.py b/aexpect/remote.py index efc9e61..1f45183 100644 --- a/aexpect/remote.py +++ b/aexpect/remote.py @@ -662,6 +662,7 @@ def remote_copy( log_function=None, transfer_timeout=600, login_timeout=300, + attempts=1, ): """ Transfer files using rsync or SCP, given a command line. @@ -677,25 +678,68 @@ def remote_copy( :param login_timeout: The maximal time duration (in seconds) to wait for each step of the login procedure (i.e. the "Are you sure" prompt or the password prompt) + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. """ - LOG.debug( - "Trying to copy with command '%s', timeout %ss", - command, - transfer_timeout, - ) - if log_filename: - output_func = log_function - output_params = (log_filename,) - else: - output_func = None - output_params = () method = "rsync" if "rsync" in command else "scp" - with Expect( - command, output_func=output_func, output_params=output_params - ) as session: - _remote_copy( - session, password_list, transfer_timeout, login_timeout, method - ) + + for attempt in range(1, attempts + 1): + try: + LOG.debug( + "Trying to copy with command '%s', timeout %ss (attempt %d/%d)", + command, + transfer_timeout, + attempt, + attempts, + ) + if log_filename: + output_func = log_function + output_params = (log_filename,) + else: + output_func = None + output_params = () + with Expect( + command, output_func=output_func, output_params=output_params + ) as session: + _remote_copy( + session, + password_list, + transfer_timeout, + login_timeout, + method, + ) + return # transfer is successful + except ( + TransferTimeoutError, + AuthenticationTimeoutError, + ExpectTimeoutError, + ) as error: + if attempt < attempts: + LOG.debug( + "Transient error on attempt %d/%d, retrying: %s", + attempt, + attempts, + error, + ) + time.sleep(1) # small delay before retry + else: + raise + except (TransferFailedError, SCPError, RsyncError) as error: + # For transfer failures, only retry on specific conditions + error_str = str(error).lower() + if "connection" in error_str or "timeout" in error_str: + if attempt < attempts: + LOG.debug( + "Connection error on attempt %d/%d, retrying: %s", + attempt, + attempts, + error, + ) + time.sleep(1) # small delay before retry + else: + raise + else: + raise def scp_to_remote( @@ -711,6 +755,7 @@ def scp_to_remote( log_function=None, timeout=600, interface=None, + attempts=1, ): """ Copy files to a remote host (guest) through scp. @@ -729,6 +774,8 @@ def scp_to_remote( to complete. :param interface: The interface the neighbours attach to (only use when using ipv6 linklocal address). + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. """ if limit: limit = f"-l {limit}" @@ -753,7 +800,12 @@ def scp_to_remote( ) password_list = [password] return remote_copy( - command, password_list, log_filename, log_function, timeout + command, + password_list, + log_filename, + log_function, + timeout, + attempts=attempts, ) @@ -770,6 +822,7 @@ def scp_from_remote( log_function=None, timeout=600, interface=None, + attempts=1, ): """ Copy files from a remote host (guest). @@ -788,6 +841,8 @@ def scp_from_remote( to complete. :param interface: The interface the neighbours attach to (only use when using ipv6 linklocal address). + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. """ if limit: limit = f"-l {limit}" @@ -810,7 +865,14 @@ def scp_from_remote( rf"{shlex.quote(local_path)}" ) password_list = [password] - remote_copy(command, password_list, log_filename, log_function, timeout) + remote_copy( + command, + password_list, + log_filename, + log_function, + timeout, + attempts=attempts, + ) def scp_between_remotes( @@ -830,6 +892,7 @@ def scp_between_remotes( timeout=600, src_inter=None, dst_inter=None, + attempts=1, ): """ Copy files from a remote host (guest) to another remote host (guest). @@ -851,6 +914,8 @@ def scp_between_remotes( to complete. :param src_inter: The interface on local that the src neighbour attached :param dst_inter: The interface on the src that the dst neighbour attached + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. :return: True on success and False on failure. """ @@ -883,7 +948,12 @@ def scp_between_remotes( ) password_list = [s_passwd, d_passwd] return remote_copy( - command, password_list, log_filename, log_function, timeout + command, + password_list, + log_filename, + log_function, + timeout, + attempts=attempts, ) @@ -900,6 +970,7 @@ def rsync_to_remote( log_function=None, timeout=600, interface=None, + attempts=1, ): """ Copy files to a remote host (guest) through rsync. @@ -918,6 +989,8 @@ def rsync_to_remote( to complete. :param interface: The interface the neighbours attach to (only use when using ipv6 linklocal address). + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. :raise: Whatever remote_rsync() raises """ if limit: @@ -941,7 +1014,12 @@ def rsync_to_remote( ) password_list = [password] return remote_copy( - command, password_list, log_filename, log_function, timeout + command, + password_list, + log_filename, + log_function, + timeout, + attempts=attempts, ) @@ -958,6 +1036,7 @@ def rsync_from_remote( log_function=None, timeout=600, interface=None, + attempts=1, ): """ Copy files from a remote host (guest) through rsync. @@ -976,6 +1055,8 @@ def rsync_from_remote( to complete. :param interface: The interface the neighbours attach to (only use when using ipv6 linklocal address). + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. :raise: Whatever remote_rsync() raises """ if limit: @@ -997,7 +1078,14 @@ def rsync_from_remote( f"{username}@{host}:{quote_path(remote_path)} {shlex.quote(local_path)}" ) password_list = [password] - remote_copy(command, password_list, log_filename, log_function, timeout) + remote_copy( + command, + password_list, + log_filename, + log_function, + timeout, + attempts=attempts, + ) # noinspection PyBroadException @@ -1020,11 +1108,13 @@ def nc_copy_between_remotes( s_session=None, d_session=None, file_transfer_timeout=600, + attempts=1, ): """ Copy files from guest to guest using netcat. This method only supports linux guest OS. + Caller-provided sessions are left open. :param src: Hostname or IP address of source :param dst: Hostname or IP address of destination @@ -1045,59 +1135,95 @@ def nc_copy_between_remotes( :param d_session: A shell session object for dst or None. :param check_sum: Whether to run checksum for the operation. :param file_transfer_timeout: Timeout for file transfer. - + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. :return: True on success and False on failure. """ - check_string = "NCFT" - if not s_session: - s_session = remote_login( - c_type, src, s_port, s_name, s_passwd, c_prompt - ) - if not d_session: - d_session = remote_login( - c_type, dst, s_port, d_name, d_passwd, c_prompt - ) + close_s_session = s_session is None + close_d_session = d_session is None + for attempt in range(1, attempts + 1): + try: + check_string = "NCFT" + if s_session is None: + s_session = remote_login( + c_type, src, s_port, s_name, s_passwd, c_prompt + ) + if d_session is None: + d_session = remote_login( + c_type, dst, s_port, d_name, d_passwd, c_prompt + ) - try: - s_session.cmd(f"iptables -I INPUT -p {d_protocol} -j ACCEPT") - d_session.cmd(f"iptables -I OUTPUT -p {d_protocol} -j ACCEPT") - except Exception: # pylint: disable=W0703 - pass - - LOG.info("Transfer data using netcat from %s to %s", src, dst) - cmd = f"nc -w {timeout}" - if d_protocol == "udp": - cmd += " -u" - receive_cmd = f"echo {check_string} | {cmd} -l {d_port} > {d_path}" - d_session.sendline(receive_cmd) - send_cmd = f"{cmd} {dst} {d_port} < {s_path}" - status, output = s_session.cmd_status_output( - send_cmd, timeout=file_transfer_timeout - ) - if status: - err = f"Fail to transfer file between {src} -> {dst}." - if check_string not in output: - err += ( - "src did not receive check " - f"string {check_string} sent by dst." + try: + s_session.cmd(f"iptables -I INPUT -p {d_protocol} -j ACCEPT") + d_session.cmd(f"iptables -I OUTPUT -p {d_protocol} -j ACCEPT") + except Exception: # pylint: disable=W0703 + pass + + LOG.info( + "Transfer data using netcat from %s to %s (attempt %d/%d)", + src, + dst, + attempt, + attempts, ) - err += f"send nc command {send_cmd}, output {output}" - err += f"Receive nc command {receive_cmd}." - raise NetcatTransferFailedError(status, err) - - if check_sum: - LOG.info("md5sum cmd = md5sum %s", s_path) - output = s_session.cmd(f"md5sum {s_path}") - src_md5 = output.split()[0] - dst_md5 = d_session.cmd(f"md5sum {d_path}").split()[0] - if src_md5.strip() != dst_md5.strip(): - err_msg = ( - "Files md5sum mismatch, " - f"file {s_path} md5sum is '{src_md5}', " - f"but the file {d_path} md5sum is {dst_md5}" + cmd = f"nc -w {timeout}" + if d_protocol == "udp": + cmd += " -u" + receive_cmd = f"echo {check_string} | {cmd} -l {d_port} > {d_path}" + d_session.sendline(receive_cmd) + send_cmd = f"{cmd} {dst} {d_port} < {s_path}" + status, output = s_session.cmd_status_output( + send_cmd, timeout=file_transfer_timeout ) - raise NetcatTransferIntegrityError(err_msg) - return True + if status: + err = f"Fail to transfer file between {src} -> {dst}." + if check_string not in output: + err += ( + "src did not receive check " + f"string {check_string} sent by dst." + ) + err += f"send nc command {send_cmd}, output {output}" + err += f"Receive nc command {receive_cmd}." + raise NetcatTransferFailedError(status, err) + + if check_sum: + LOG.info("md5sum cmd = md5sum %s", s_path) + output = s_session.cmd(f"md5sum {s_path}") + src_md5 = output.split()[0] + dst_md5 = d_session.cmd(f"md5sum {d_path}").split()[0] + if src_md5.strip() != dst_md5.strip(): + err_msg = ( + "Files md5sum mismatch, " + f"file {s_path} md5sum is '{src_md5}', " + f"but the file {d_path} md5sum is {dst_md5}" + ) + raise NetcatTransferIntegrityError(err_msg) + return True + except ( + NetcatTransferTimeoutError, + NetcatTransferFailedError, + UDPError, + ) as error: + if attempt < attempts: + LOG.debug( + "Transfer failed on attempt %d/%d, retrying: %s", + attempt, + attempts, + error, + ) + time.sleep(1) # small delay before retry + else: + raise + finally: + try: + if close_s_session and s_session is not None: + s_session.close() + s_session = None + finally: + if close_d_session and d_session is not None: + d_session.close() + d_session = None + return False def udp_copy_between_remotes( @@ -1114,6 +1240,7 @@ def udp_copy_between_remotes( c_prompt="\n", d_port="9000", timeout=600, + attempts=1, ): """ Copy files from guest to guest using udp. @@ -1131,9 +1258,9 @@ def udp_copy_between_remotes( :param c_prompt: command line prompt of remote host(guest) :param d_port: the port data transfer :param timeout: data transfer timeout + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. """ - s_session = remote_login(c_type, src, s_port, s_name, s_passwd, c_prompt) - d_session = remote_login(c_type, dst, s_port, d_name, d_passwd, c_prompt) def get_abs_path(session, filename, extension): """Return file path drive+path.""" @@ -1215,23 +1342,56 @@ def stop_server(session): if server_alive(session): session.cmd_output_safe(stop_cmd) - try: - src_md5 = get_file_md5(s_session, s_path) - if not server_alive(s_session): - start_server(s_session) - start_client(d_session) - dst_md5 = get_file_md5(d_session, d_path) - if src_md5 != dst_md5: - err_msg = ( - "Files md5sum mismatch, " - f"file {s_path} md5sum is '{src_md5}', " - f"but the file {d_path} md5sum is {dst_md5}" - ) - raise UDPError(err_msg) - finally: - stop_server(s_session) - s_session.close() - d_session.close() + def close_sessions(s_session, d_session): + """Close both sessions even if closing the source session fails.""" + try: + if s_session is not None: + s_session.close() + finally: + if d_session is not None: + d_session.close() + + for attempt in range(1, attempts + 1): + s_session = None + d_session = None + try: + try: + s_session = remote_login( + c_type, src, s_port, s_name, s_passwd, c_prompt + ) + d_session = remote_login( + c_type, dst, s_port, d_name, d_passwd, c_prompt + ) + src_md5 = get_file_md5(s_session, s_path) + if not server_alive(s_session): + start_server(s_session) + start_client(d_session) + dst_md5 = get_file_md5(d_session, d_path) + if src_md5 != dst_md5: + err_msg = ( + "Files md5sum mismatch, " + f"file {s_path} md5sum is '{src_md5}', " + f"but the file {d_path} md5sum is {dst_md5}" + ) + raise UDPError(err_msg) + finally: + try: + if s_session is not None and d_session is not None: + stop_server(s_session) + finally: + close_sessions(s_session, d_session) + return # transfer is successful + except UDPError as error: + if attempt < attempts: + LOG.debug( + "UDP transfer failed on attempt %d/%d, retrying: %s", + attempt, + attempts, + error, + ) + time.sleep(1) # small delay before retry + else: + raise def login_from_session( @@ -1284,6 +1444,7 @@ def scp_to_session( log_function=None, timeout=600, interface=None, + attempts=1, ): """ Secure copy a filepath (w/o wildcard) to a remote location with the same @@ -1299,6 +1460,8 @@ def scp_to_session( :param log_function: Function to perform logging :param timeout: Timeout for the scp operation :param interface: Interface used for the transfer + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. The rest of the arguments are identical to scp_to_remote(). """ @@ -1315,6 +1478,7 @@ def scp_to_session( log_function, timeout, interface, + attempts, ) @@ -1328,6 +1492,7 @@ def scp_from_session( log_function=None, timeout=600, interface=None, + attempts=1, ): """ Secure copy a filepath (w/o wildcard) from a remote location with the same @@ -1343,6 +1508,8 @@ def scp_from_session( :param log_function: Function to perform logging :param timeout: Timeout for the scp operation :param interface: Interface used for the transfer + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. The rest of the arguments are identical to scp_from_remote(). """ @@ -1359,6 +1526,7 @@ def scp_from_session( log_function, timeout, interface, + attempts, ) @@ -1406,6 +1574,7 @@ def copy_files_to( timeout=600, interface=None, filesize=None, # pylint: disable=unused-argument + attempts=1, ): """ Copy files to a remote host (guest) using the selected client. @@ -1427,6 +1596,8 @@ def copy_files_to( :param interface: The interface the neighbours attach to (only use when using ipv6 linklocal address.) :param filesize: size of file will be transferred + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. """ if client == "scp": scp_to_remote( @@ -1442,6 +1613,7 @@ def copy_files_to( log_function, timeout, interface=interface, + attempts=attempts, ) elif client == "rsync": rsync_to_remote( @@ -1457,6 +1629,7 @@ def copy_files_to( log_function, timeout, interface=interface, + attempts=attempts, ) elif client == "rss": log_func = None @@ -1464,9 +1637,27 @@ def copy_files_to( log_func = LOG.debug if interface: address = f"{address}%{interface}" - fdclient = rss_client.FileUploadClient(address, port, log_func) - fdclient.upload(local_path, remote_path, timeout) - fdclient.close() + for attempt in range(1, attempts + 1): + try: + fdclient = rss_client.FileUploadClient(address, port, log_func) + fdclient.upload(local_path, remote_path, timeout) + fdclient.close() + return # transfer is successful + except ( + rss_client.FileTransferConnectError, + rss_client.FileTransferTimeoutError, + rss_client.FileTransferSocketError, + ) as error: + if attempt < attempts: + LOG.debug( + "RSS upload failed on attempt %d/%d, retrying: %s", + attempt, + attempts, + error, + ) + time.sleep(1) # small delay before retry + else: + raise else: raise TransferBadClientError(client) @@ -1489,6 +1680,7 @@ def copy_files_from( timeout=600, interface=None, filesize=None, # pylint: disable=unused-argument + attempts=1, ): """ Copy files from a remote host (guest) using the selected client. @@ -1510,6 +1702,8 @@ def copy_files_from( :param interface: The interface the neighbours attach to (only use when using ipv6 linklocal address.) :param filesize: size of file will be transferred + :param attempts: Number of attempts to make to deal with transient errors like + timeouts and connection issues. """ if client == "scp": scp_from_remote( @@ -1525,6 +1719,7 @@ def copy_files_from( log_function, timeout, interface=interface, + attempts=attempts, ) elif client == "rsync": rsync_from_remote( @@ -1540,6 +1735,7 @@ def copy_files_from( log_function, timeout, interface=interface, + attempts=attempts, ) elif client == "rss": log_func = None @@ -1547,8 +1743,28 @@ def copy_files_from( log_func = LOG.debug if interface: address = f"{address}%{interface}" - fdclient = rss_client.FileDownloadClient(address, port, log_func) - fdclient.download(remote_path, local_path, timeout) - fdclient.close() + for attempt in range(1, attempts + 1): + try: + fdclient = rss_client.FileDownloadClient( + address, port, log_func + ) + fdclient.download(remote_path, local_path, timeout) + fdclient.close() + return # transfer is successful + except ( + rss_client.FileTransferConnectError, + rss_client.FileTransferTimeoutError, + rss_client.FileTransferSocketError, + ) as error: + if attempt < attempts: + LOG.debug( + "RSS download failed on attempt %d/%d, retrying: %s", + attempt, + attempts, + error, + ) + time.sleep(1) # small delay before retry + else: + raise else: raise TransferBadClientError(client) diff --git a/tests/test_remote.py b/tests/test_remote.py index a590360..2ce96cd 100644 --- a/tests/test_remote.py +++ b/tests/test_remote.py @@ -32,6 +32,9 @@ def setUp(self): ) session_patch.start() self.addCleanup(session_patch.stop) + expect_patch = mock.patch("aexpect.remote.Expect") + self.expect = expect_patch.start() + self.addCleanup(expect_patch.stop) def test_handle_prompts(self): output = remote.handle_prompts( @@ -59,6 +62,22 @@ def test_wait_for_login(self): " -o PreferredAuthentications=password user@127.0.0.1", ) + @mock.patch("aexpect.remote._remote_copy") + def test_remote_copy(self, mock_remote_copy): + remote.remote_copy("cp a b", ["pass"]) + mock_remote_copy.assert_called_once_with( + mock.ANY, + ["pass"], + 600, + 300, + "scp", + ) + self.expect.assert_called_once_with( + r"cp a b", + output_func=None, + output_params=(), + ) + @mock.patch("aexpect.remote._remote_copy") def test_scp_to_remote(self, mock_remote_copy): remote.scp_to_remote( @@ -67,9 +86,10 @@ def test_scp_to_remote(self, mock_remote_copy): mock_remote_copy.assert_called_once_with( mock.ANY, ["pass"], 600, 300, "scp" ) - self.assertEqual( - mock_remote_copy.call_args[0][0].command, + self.expect.assert_called_once_with( r"scp -r -v -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o PreferredAuthentications=password -P 22 /local/path user@\[127.0.0.1\]:/remote/path", + output_func=None, + output_params=(), ) @mock.patch("aexpect.remote._remote_copy") @@ -80,22 +100,10 @@ def test_scp_from_remote(self, mock_remote_copy): mock_remote_copy.assert_called_once_with( mock.ANY, ["pass"], 600, 300, "scp" ) - self.assertEqual( - mock_remote_copy.call_args[0][0].command, + self.expect.assert_called_once_with( r"scp -r -v -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o PreferredAuthentications=password -P 22 user@\[127.0.0.1\]:/remote/path /local/path", - ) - - @mock.patch("aexpect.remote._remote_copy") - def test_rsync_to_remote(self, mock_remote_copy): - remote.rsync_to_remote( - "127.0.0.1", 22, "user", "pass", "/local/path", "/remote/path" - ) - mock_remote_copy.assert_called_once_with( - mock.ANY, ["pass"], 600, 300, "rsync" - ) - self.assertEqual( - mock_remote_copy.call_args[0][0].command, - r"rsync -r -avz -e 'ssh -Tp 22 -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no' /local/path user@127.0.0.1:/remote/path", + output_func=None, + output_params=(), ) @mock.patch("aexpect.remote._remote_copy") @@ -114,9 +122,24 @@ def test_scp_between_remotes(self, mock_remote_copy): mock_remote_copy.assert_called_once_with( mock.ANY, ["src_pass", "dst_pass"], 600, 300, "scp" ) - self.assertEqual( - mock_remote_copy.call_args[0][0].command, + self.expect.assert_called_once_with( r"scp -r -v -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o PreferredAuthentications=password -P 22 src_user@\[src_host\]:/src/path dst_user@\[dst_host\]:/dst/path", + output_func=None, + output_params=(), + ) + + @mock.patch("aexpect.remote._remote_copy") + def test_rsync_to_remote(self, mock_remote_copy): + remote.rsync_to_remote( + "127.0.0.1", 22, "user", "pass", "/local/path", "/remote/path" + ) + mock_remote_copy.assert_called_once_with( + mock.ANY, ["pass"], 600, 300, "rsync" + ) + self.expect.assert_called_once_with( + r"rsync -r -avz -e 'ssh -Tp 22 -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no' /local/path user@127.0.0.1:/remote/path", + output_func=None, + output_params=(), ) @mock.patch("aexpect.remote._remote_copy") @@ -127,7 +150,444 @@ def test_rsync_from_remote(self, mock_remote_copy): mock_remote_copy.assert_called_once_with( mock.ANY, ["pass"], 600, 300, "rsync" ) - self.assertEqual( - mock_remote_copy.call_args[0][0].command, + self.expect.assert_called_once_with( r"rsync -r -avz -e 'ssh -Tp 22 -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no' user@127.0.0.1:/remote/path /local/path", + output_func=None, + output_params=(), + ) + + +class TestRemoteCopyRetry(unittest.TestCase): + + def setUp(self): + copy_patch = mock.patch("aexpect.remote._remote_copy") + self.copy = copy_patch.start() + self.addCleanup(copy_patch.stop) + expect_patch = mock.patch("aexpect.remote.Expect") + self.expect = expect_patch.start() + self.addCleanup(expect_patch.stop) + sleep_patch = mock.patch("aexpect.remote.time.sleep") + self.sleep = sleep_patch.start() + self.addCleanup(sleep_patch.stop) + self.retry_errors = ( + remote.TransferTimeoutError("Transfer stalled"), + remote.AuthenticationTimeoutError("Login stalled"), + remote.ExpectTimeoutError(["password"], "No response"), + remote.TransferFailedError(1, "Connection reset"), + remote.SCPError("Copy failed", "CONNECTION lost"), + remote.RsyncError("Copy failed", "TIMEOUT"), + ) + + def assert_attempts(self, attempts): + self.assertEqual(self.copy.call_count, attempts) + self.assertEqual(self.expect.call_count, attempts) + self.assertEqual( + self.expect.return_value.__exit__.call_count, attempts + ) + self.assertEqual( + self.sleep.call_args_list, [mock.call(1)] * (attempts - 1) + ) + + def test_success_does_not_retry(self): + remote.remote_copy("scp source destination", ["pass"], attempts=3) + self.assert_attempts(1) + + def test_default_attempt_does_not_retry(self): + self.copy.side_effect = remote.TransferTimeoutError("Transfer stalled") + with self.assertRaises(remote.TransferTimeoutError): + remote.remote_copy("scp source destination", ["pass"]) + self.assert_attempts(1) + + def test_transient_errors_retry_until_success(self): + for error in self.retry_errors: + with self.subTest(error=error): + for mocked in (self.copy, self.expect, self.sleep): + mocked.reset_mock() + self.copy.side_effect = [error, None] + remote.remote_copy( + "scp source destination", ["pass"], attempts=3 + ) + self.assert_attempts(2) + + def test_exhausted_attempts_raise_last_error(self): + for error in self.retry_errors: + for attempts in (1, 3): + with self.subTest(error=error, attempts=attempts): + for mocked in (self.copy, self.expect, self.sleep): + mocked.reset_mock() + self.copy.side_effect = [ + remote.TransferTimeoutError("Earlier failure") + ] * (attempts - 1) + [error] + with self.assertRaises(type(error)) as raised: + remote.remote_copy( + "scp source destination", + ["pass"], + attempts=attempts, + ) + self.assertIs(raised.exception, error) + self.assert_attempts(attempts) + + def test_permanent_errors_do_not_retry(self): + for error in ( + remote.TransferFailedError(1, "Permission denied"), + remote.SCPError("Copy failed", "No such file"), + remote.RsyncError("Copy failed", "No such file"), + remote.AuthenticationError("Login failed", "Permission denied"), + ValueError("Invalid transfer"), + ): + with self.subTest(error=error): + for mocked in (self.copy, self.expect, self.sleep): + mocked.reset_mock() + self.copy.side_effect = error + with self.assertRaises(type(error)) as raised: + remote.remote_copy( + "scp source destination", ["pass"], attempts=3 + ) + self.assertIs(raised.exception, error) + self.assert_attempts(1) + + def test_wrappers_preserve_attempts(self): + args = ("host", 22, "user", "pass", "/source", "/destination") + session = mock.Mock( + host="host", port=22, username="user", password="pass" ) + cases = [ + (remote.scp_to_remote, args), + (remote.scp_from_remote, args), + (remote.rsync_to_remote, args), + (remote.rsync_from_remote, args), + ( + remote.scp_between_remotes, + ( + "src", + "dst", + 22, + "pass", + "pass", + "user", + "user", + "/source", + "/destination", + ), + ), + (remote.scp_to_session, (session, "/source", "/destination")), + (remote.scp_from_session, (session, "/source", "/destination")), + ] + for client in ("scp", "rsync"): + for method in (remote.copy_files_to, remote.copy_files_from): + cases.append( + ( + method, + ( + "host", + client, + "user", + "pass", + 22, + "/source", + "/destination", + ), + ) + ) + for method, args in cases: + with self.subTest(method=method, args=args): + for mocked in (self.copy, self.expect, self.sleep): + mocked.reset_mock() + self.copy.side_effect = [ + remote.TransferTimeoutError("Transfer stalled"), + remote.TransferTimeoutError("Transfer stalled again"), + None, + ] + method(*args, attempts=3) + self.assert_attempts(3) + + +class TestRSSCopyRetry(unittest.TestCase): + + def setUp(self): + self.args = ( + "host", + "rss", + "user", + "pass", + 22, + "/source", + "/destination", + ) + self.cases = ( + (remote.copy_files_to, "FileUploadClient", "upload"), + (remote.copy_files_from, "FileDownloadClient", "download"), + ) + self.retry_errors = ( + ( + True, + remote.rss_client.FileTransferConnectError( + "Connection failed" + ), + ), + ( + False, + remote.rss_client.FileTransferTimeoutError("Transfer stalled"), + ), + ( + False, + remote.rss_client.FileTransferSocketError("Connection lost"), + ), + ) + sleep_patch = mock.patch("aexpect.remote.time.sleep") + self.sleep = sleep_patch.start() + self.addCleanup(sleep_patch.stop) + + def test_connection_and_transfer_retry_until_success(self): + for method, client_class, operation in self.cases: + for connection_failure, error in self.retry_errors: + with self.subTest(client=client_class, error=error): + self.sleep.reset_mock() + first, second = mock.Mock(), mock.Mock() + getattr(first, operation).side_effect = error + with mock.patch.object( + remote.rss_client, client_class + ) as client: + client.side_effect = [ + error if connection_failure else first, + second, + ] + method(*self.args, attempts=3) + self.assertEqual(client.call_count, 2) + getattr(second, operation).assert_called_once_with( + "/source", "/destination", 600 + ) + second.close.assert_called_once_with() + self.sleep.assert_called_once_with(1) + + def test_connection_and_transfer_attempts_exhausted(self): + for method, client_class, operation in self.cases: + for connection_failure, error in self.retry_errors: + for attempts in (1, 3): + with self.subTest( + client=client_class, + error=error, + attempts=attempts, + ): + self.sleep.reset_mock() + with mock.patch.object( + remote.rss_client, client_class + ) as client: + if connection_failure: + client.side_effect = error + else: + getattr( + client.return_value, operation + ).side_effect = error + with self.assertRaises(type(error)) as raised: + method(*self.args, attempts=attempts) + self.assertIs(raised.exception, error) + self.assertEqual(client.call_count, attempts) + self.assertEqual( + self.sleep.call_args_list, + [mock.call(1)] * (attempts - 1), + ) + + def test_other_errors_do_not_retry(self): + for method, client_class, operation in self.cases: + for error in ( + remote.rss_client.FileTransferError("Transfer failed"), + remote.rss_client.FileTransferNotFoundError("No such file"), + remote.rss_client.FileTransferProtocolError("Invalid message"), + remote.rss_client.FileTransferServerError("Permission denied"), + TypeError("Invalid argument"), + ValueError("Invalid value"), + ): + with self.subTest(client=client_class, error=error): + self.sleep.reset_mock() + with mock.patch.object( + remote.rss_client, client_class + ) as client: + transfer = getattr(client.return_value, operation) + transfer.side_effect = error + with self.assertRaises(type(error)) as raised: + method(*self.args, attempts=3) + self.assertIs(raised.exception, error) + client.assert_called_once_with("host", 22, None) + transfer.assert_called_once_with( + "/source", "/destination", 600 + ) + self.sleep.assert_not_called() + + +class TestTransferSessionCleanup(unittest.TestCase): + + def setUp(self): + self.args = ( + "src", + "dst", + 22, + "pass", + "pass", + "user", + "user", + "/src/path", + "/dst/path", + ) + self.s_session = mock.Mock(spec=RemoteSession) + self.d_session = mock.Mock(spec=RemoteSession) + self.s_session.cmd_status_output.return_value = (0, "NCFT") + self.s_session.cmd.return_value = "abc" + self.d_session.cmd.return_value = "abc" + self.s_session.cmd_output.return_value = "abc /src/path\nsendfile" + self.d_session.cmd_output.return_value = "abc /dst/path" + login_patch = mock.patch("aexpect.remote.remote_login") + self.login = login_patch.start() + self.addCleanup(login_patch.stop) + self.login.side_effect = [self.s_session, self.d_session] + sleep_patch = mock.patch("aexpect.remote.time.sleep") + self.sleep = sleep_patch.start() + self.addCleanup(sleep_patch.stop) + + def test_login_failure_closes_created_sessions(self): + error = remote.LoginError("Login failed") + for method in ( + remote.nc_copy_between_remotes, + remote.udp_copy_between_remotes, + ): + for source_connected in (False, True): + with self.subTest(method=method, source=source_connected): + self.s_session.reset_mock() + self.login.side_effect = ( + [self.s_session, error] + if source_connected + else [error] + ) + with self.assertRaises(remote.LoginError) as raised: + method(*self.args) + self.assertIs(raised.exception, error) + self.assertEqual( + self.s_session.close.call_count, int(source_connected) + ) + + def test_nc_closes_only_created_sessions(self): + for supply_source, supply_destination in ( + (False, False), + (True, False), + (False, True), + (True, True), + ): + for status in (0, 1): + with self.subTest( + source=supply_source, + destination=supply_destination, + status=status, + ): + self.s_session.reset_mock() + self.d_session.reset_mock() + self.login.reset_mock() + self.login.side_effect = [ + session + for session, supplied in ( + (self.s_session, supply_source), + (self.d_session, supply_destination), + ) + if not supplied + ] + self.s_session.cmd_status_output.return_value = ( + status, + "NCFT", + ) + kwargs = { + "s_session": self.s_session if supply_source else None, + "d_session": ( + self.d_session if supply_destination else None + ), + } + if status: + with self.assertRaises( + remote.NetcatTransferFailedError + ): + remote.nc_copy_between_remotes( + *self.args, **kwargs + ) + else: + self.assertTrue( + remote.nc_copy_between_remotes( + *self.args, **kwargs + ) + ) + self.assertEqual( + self.s_session.close.call_count, int(not supply_source) + ) + self.assertEqual( + self.d_session.close.call_count, + int(not supply_destination), + ) + self.assertEqual( + self.login.call_count, + int(not supply_source) + int(not supply_destination), + ) + + def test_nc_retry_preserves_supplied_session(self): + replacement = mock.Mock(spec=RemoteSession) + self.login.side_effect = [self.d_session, replacement] + self.s_session.cmd_status_output.side_effect = [ + (1, "NCFT"), + (0, "NCFT"), + ] + self.assertTrue( + remote.nc_copy_between_remotes( + *self.args, + s_session=self.s_session, + check_sum=False, + attempts=2, + ) + ) + self.s_session.close.assert_not_called() + self.d_session.close.assert_called_once_with() + replacement.close.assert_called_once_with() + self.assertEqual(self.login.call_count, 2) + self.sleep.assert_called_once_with(1) + + def test_udp_success_closes_sessions(self): + remote.udp_copy_between_remotes(*self.args) + self.s_session.close.assert_called_once_with() + self.d_session.close.assert_called_once_with() + self.s_session.cmd_output_safe.assert_called_once_with( + "killall sendfile" + ) + + def test_udp_stop_failure_closes_sessions(self): + error = remote.UDPError("Cannot stop server") + self.s_session.cmd_output_safe.side_effect = error + with self.assertRaises(remote.UDPError) as raised: + remote.udp_copy_between_remotes(*self.args) + self.assertIs(raised.exception, error) + self.s_session.close.assert_called_once_with() + self.d_session.close.assert_called_once_with() + + def test_udp_retry_does_not_reuse_previous_sessions(self): + error = remote.LoginError("Login failed") + self.login.side_effect = [self.s_session, self.d_session, error] + self.d_session.cmd_output_safe.side_effect = remote.UDPError( + "Transfer failed" + ) + with self.assertRaises(remote.LoginError) as raised: + remote.udp_copy_between_remotes(*self.args, attempts=2) + self.assertIs(raised.exception, error) + self.s_session.close.assert_called_once_with() + self.d_session.close.assert_called_once_with() + self.sleep.assert_called_once_with(1) + + def test_source_close_failure_still_closes_destination(self): + error = RuntimeError("Cannot close source") + self.s_session.close.side_effect = error + for method in ( + remote.nc_copy_between_remotes, + remote.udp_copy_between_remotes, + ): + with self.subTest(method=method): + self.s_session.reset_mock() + self.d_session.reset_mock() + self.login.side_effect = [self.s_session, self.d_session] + with self.assertRaises(RuntimeError) as raised: + method(*self.args) + self.assertIs(raised.exception, error) + self.s_session.close.assert_called_once_with() + self.d_session.close.assert_called_once_with()