From c876fabd0726b33a34f72a1af7588da0360ae37f Mon Sep 17 00:00:00 2001
From: zhangsong <zhangsong@cmss.chinamobile.com>
Date: Thu, 12 May 2016 09:37:59 +0800
Subject: [PATCH] Add timeout option to ssh_execute

Add an optional timeout parameter to ssh_execute() method of
oslo_concurrency.processutils. Timeout is None means infinity,
which is the default.

Change-Id: I45532799ecd23ff8c5572036d06725e71965b2b5
Closes-Bug: #1580826
---
 oslo_concurrency/processutils.py                 |  5 +++--
 oslo_concurrency/tests/unit/test_processutils.py | 11 ++++++++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/oslo_concurrency/processutils.py b/oslo_concurrency/processutils.py
index 1f46221..808ae2a 100644
--- a/oslo_concurrency/processutils.py
+++ b/oslo_concurrency/processutils.py
@@ -466,7 +466,7 @@ def trycmd(*args, **kwargs):
 
 def ssh_execute(ssh, cmd, process_input=None,
                 addl_env=None, check_exit_code=True,
-                binary=False):
+                binary=False, timeout=None):
     """Run a command through SSH.
 
     .. versionchanged:: 1.9
@@ -481,7 +481,8 @@ def ssh_execute(ssh, cmd, process_input=None,
         # This is (probably) fixable if we need it...
         raise InvalidArgumentError(_('process_input not supported over SSH'))
 
-    stdin_stream, stdout_stream, stderr_stream = ssh.exec_command(cmd)
+    stdin_stream, stdout_stream, stderr_stream = ssh.exec_command(
+        cmd, timeout=timeout)
     channel = stdout_stream.channel
 
     # NOTE(justinsb): This seems suspicious...
diff --git a/oslo_concurrency/tests/unit/test_processutils.py b/oslo_concurrency/tests/unit/test_processutils.py
index 5c744f5..8454dbf 100644
--- a/oslo_concurrency/tests/unit/test_processutils.py
+++ b/oslo_concurrency/tests/unit/test_processutils.py
@@ -21,6 +21,7 @@ import multiprocessing
 import os
 import pickle
 import resource
+import socket
 import stat
 import subprocess
 import sys
@@ -599,7 +600,9 @@ class FakeSshConnection(object):
         self.out = out
         self.err = err
 
-    def exec_command(self, cmd):
+    def exec_command(self, cmd, timeout=None):
+        if timeout:
+            raise socket.timeout()
         stdout = FakeSshStream(self.out)
         stdout.setup_channel(self.rc)
         return (six.BytesIO(),
@@ -618,6 +621,12 @@ class SshExecuteTestCase(test_base.BaseTestCase):
                           processutils.ssh_execute,
                           None, 'ls', process_input='important')
 
+    def test_timeout_error(self):
+        self.assertRaises(socket.timeout,
+                          processutils.ssh_execute,
+                          FakeSshConnection(0), 'ls',
+                          timeout=10)
+
     def test_works(self):
         out, err = processutils.ssh_execute(FakeSshConnection(0), 'ls')
         self.assertEqual('stdout', out)