Merge lp:~abentley/workspace-runner/scp-retry into lp:workspace-runner

[Aaron Bentley]

Retry scp on retcode 1.

This branch updates the workspace runner to retry scp when it exits with status 1.

[Curtis Hovey (sinzui)]

Thank you.

review: Approve (code)

=== modified file 'workspace_runner/__init__.py'
--- workspace_runner/__init__.py	2015-06-25 15:06:58 +0000
+++ workspace_runner/__init__.py	2015-06-25 17:06:42 +0000
@@ -38,23 +38,28 @@
         return cls(ssh_connection, tempdir)
 
 
-def retry_ssh(func):
-    """Decorator that retries a function if it looks like SSH failed.
+def retry_ssh(retcodes=(255,)):
+    """Return a decorator that retries a function if it looks like SSH failed.
 
     SSH connection failures are a CalledProcessError with returncode 255.
+
+    Custom retcodes can be specified for operations that have different
+    failures, like scp, which returns 1.
     """
-    def decorated(*args, **kwargs):
-        for x in range(3):
-            if x != 0:
-                logging.info('Retrying due to ssh failure.')
-            try:
-                return func(*args, **kwargs)
-            except subprocess.CalledProcessError as e:
-                if e.returncode != 255:
-                    raise
-        else:
-            raise
-    return decorated
+    def decorator(func):
+        def decorated(*args, **kwargs):
+            for x in range(3):
+                if x != 0:
+                    logging.info('Retrying due to ssh failure.')
+                try:
+                    return func(*args, **kwargs)
+                except subprocess.CalledProcessError as e:
+                    if e.returncode not in retcodes:
+                        raise
+            else:
+                raise
+        return decorated
+    return decorator
 
 
 class SSHConnection:
@@ -81,7 +86,7 @@
         command = ['ssh'] + self.get_ssh_options() + [self.host] + args
         return command
 
-    @retry_ssh
+    @retry_ssh()
     def call_with_input(self, args, stdin):
         """Run a command remotely, accepting stdin."""
         remote_command = ' '.join(quote(x) for x in args)
@@ -96,7 +101,7 @@
                 retcode, command, output=output)
         return output[0]
 
-    @retry_ssh
+    @retry_ssh((1, 255))
     def scp_to_remote(self, sources, dest_path):
         """Copy files to the specified path on the remote host."""
         dest_loc = '{}:{}/'.format(self.host, quote(dest_path))
@@ -131,7 +136,7 @@
             shutil.rmtree(sys.argv[1])
             """), [self.workspace])
 
-    @retry_ssh
+    @retry_ssh()
     def run(self, args, out_file):
         """Run a command in the workspace."""
         command = self.ssh_connection.get_ssh_cmd(
 
=== modified file 'workspace_runner/tests/__init__.py'
--- workspace_runner/tests/__init__.py	2015-06-25 15:06:58 +0000
+++ workspace_runner/tests/__init__.py	2015-06-25 17:06:42 +0000
@@ -130,7 +130,7 @@
     def test_retry_until_failure(self):
         attempts = []
 
-        @retry_ssh
+        @retry_ssh()
         def fail_repeatedly():
             attempts.append('attempt')
             raise subprocess.CalledProcessError(255, [])
@@ -142,7 +142,7 @@
     def test_immediate_success(self):
         attempts = []
 
-        @retry_ssh
+        @retry_ssh()
         def succeed():
             attempts.append('success')
             return 'success'
@@ -153,7 +153,7 @@
     def test_no_retry_non_ssh(self):
         attempts = []
 
-        @retry_ssh
+        @retry_ssh()
         def fail_repeatedly():
             attempts.append('attempt')
             raise subprocess.CalledProcessError(254, [])
@@ -162,10 +162,22 @@
             fail_repeatedly()
         self.assertEqual(len(attempts), 1)
 
+    def test_custom_retcodes(self):
+        attempts = []
+
+        @retry_ssh((254,))
+        def fail_repeatedly():
+            attempts.append('attempt')
+            raise subprocess.CalledProcessError(254, [])
+
+        with self.assertRaises(subprocess.CalledProcessError):
+            fail_repeatedly()
+        self.assertEqual(len(attempts), 3)
+
     def test_nearly_fail(self):
         attempts = []
 
-        @retry_ssh
+        @retry_ssh()
         def nearly_fail():
             attempts.append('attempt')
             if len(attempts) == 3: