Allow to specify timeout for commands in shotgun

Added new global setting DEFAULT_TIMEOUT for setting
default timeout (10s) for executing command.

It's also possible to specify timeout for all commands in
snapshot by using 'timeout' option in snapshot config file.

Additionally it's also possible to set command-specific timeout
by using 'timeout' option in command config.

Change-Id: If896485ea3b20ddde7bfc955a0a12b1faef98ea8
Closes-Bug: #1441954

shotgun/config.py

@@ -62,3 +62,8 @@ class Config(object):
                 for object_ in properties.get("objects", []):
                     object_["host"] = host
                     yield object_
+
+    @property
+    def timeout(self):
+        """Timeout for executing commands."""
+        return self.data.get("timeout", settings.DEFAULT_TIMEOUT)

shotgun/driver.py

@@ -45,6 +45,7 @@ class CommandOut(object):
 
 
 class Driver(object):
+
     @classmethod
     def getDriver(cls, data, conf):
         driver_type = data["type"]
@@ -65,6 +66,7 @@ class Driver(object):
         self.ssh_key = self.data.get("host", {}).get("ssh-key")
         self.local = utils.is_local(self.host)
         self.conf = conf
+        self.timeout = self.data.get("timeout", self.conf.timeout)
 
     def snapshot(self):
         raise NotImplementedError
@@ -79,7 +81,7 @@ class Driver(object):
                     host_string=self.host,      # destination host
                     key_filename=self.ssh_key,  # a path to ssh key
                     timeout=2,                  # a network connection timeout
-                    command_timeout=10,         # a command execution timeout
+                    command_timeout=self.timeout,  # command execution timeout
                     warn_only=True,             # don't exit on error
                     abort_on_prompts=True,      # non-interactive mode
                 ):
@@ -302,6 +304,7 @@ class XmlRpc(Driver):
 
 
 class Command(Driver):
+
     def __init__(self, data, conf):
         super(Command, self).__init__(data, conf)
         self.cmdname = self.data["command"]

shotgun/settings.py

@@ -17,3 +17,4 @@ LASTDUMP = "/tmp/snapshot_last"
 TIMESTAMP = True
 COMPRESSION_LEVEL = 3
 LOG_FILE = "/var/log/shotgun.log"
+DEFAULT_TIMEOUT = 10

shotgun/test/base.py

@@ -12,11 +12,7 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
-try:
+from unittest2.case import TestCase
-    from unittest.case import TestCase
-except ImportError:
-    # Runing unit-tests in production environment
-    from unittest2.case import TestCase
 
 
 class BaseTestCase(TestCase):

shotgun/test/test_config.py

@@ -12,10 +12,10 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
-from mock import patch
-import re
 import time
 
+import mock
+
 from shotgun.config import Config
 from shotgun.test import base
 
@@ -24,7 +24,7 @@ class TestConfig(base.BaseTestCase):
 
     def test_timestamp(self):
         t = time.localtime()
-        with patch('shotgun.config.time') as MockedTime:
+        with mock.patch('shotgun.config.time') as MockedTime:
             MockedTime.localtime.return_value = t
             MockedTime.strftime.side_effect = time.strftime
             conf = Config({})
@@ -39,10 +39,20 @@ class TestConfig(base.BaseTestCase):
             "target": "/tmp/sample",
             "timestamp": True
         })
-        assert bool(
+        self.assertRegex(
-            re.search(
+            conf.target,
-                ur"\/tmp\/sample\-[\d]{4}\-[\d]{2}\-[\d]{2}_"
+            ur"\/tmp\/sample\-[\d]{4}\-[\d]{2}\-[\d]{2}_"
-                "([\d]{2}\-){2}[\d]{2}",
+            "([\d]{2}\-){2}[\d]{2}",
-                conf.target
-            )
         )
+
+    @mock.patch('shotgun.config.settings')
+    def test_timeout(self, m_settings):
+        conf = Config({})
+        self.assertIs(conf.timeout, m_settings.DEFAULT_TIMEOUT)
+
+    def test_pass_default_timeout(self):
+        timeout = 1345
+        conf = Config({
+            'timeout': timeout,
+        })
+        self.assertEqual(conf.timeout, timeout)

shotgun/test/test_driver.py

@@ -14,14 +14,13 @@
 
 import fnmatch
 import os
+import random
 import sys
 
 import fabric
 import mock
 
-import shotgun.config
+import shotgun
-import shotgun.driver
-import shotgun.settings
 from shotgun.test import base
 
 
@@ -63,17 +62,39 @@ class TestDriver(base.BaseTestCase):
 
         command = "COMMAND"
 
+        conf = mock.Mock()
         driver = shotgun.driver.Driver(
-            {"host": {"address": "remote_host"}}, None)
+            {"host": {"address": "remote_host"}}, conf)
         result = driver.command(command)
 
-        shotgun.driver.fabric.api.run.assert_called_with(
+        mfabrun.assert_called_with(
             command, stdout=mock.ANY)
-        shotgun.driver.fabric.api.settings.assert_called_with(
+        mfabset.assert_called_with(
-            host_string="remote_host", timeout=2, command_timeout=10,
+            host_string="remote_host",
-            warn_only=True, key_filename=None, abort_on_prompts=True)
+            timeout=2,
+            command_timeout=driver.timeout,
+            warn_only=True,
+            key_filename=None,
+            abort_on_prompts=True)
         self.assertEqual(result, out)
 
+    @mock.patch('shotgun.driver.fabric.api.run')
+    @mock.patch('shotgun.driver.fabric.api.settings')
+    def test_fabric_use_timout_from_driver(self, mfabset, _):
+        timeout = random.randint(1, 100)
+        conf = mock.Mock()
+        driver = shotgun.driver.Driver(
+            {"host": {"address": "remote_host"}}, conf)
+        driver.timeout = timeout
+        driver.command("COMMAND")
+        mfabset.assert_called_with(
+            host_string=mock.ANY,
+            timeout=mock.ANY,
+            command_timeout=timeout,
+            warn_only=mock.ANY,
+            key_filename=mock.ANY,
+            abort_on_prompts=mock.ANY)
+
     @mock.patch('shotgun.driver.utils.execute')
     def test_driver_local_command(self, mexecute):
         mexecute.return_value = ("RETURN_CODE", "STDOUT", "STDERR")
@@ -84,7 +105,8 @@ class TestDriver(base.BaseTestCase):
         out.return_code = "RETURN_CODE"
 
         command = "COMMAND"
-        driver = shotgun.driver.Driver({}, None)
+        conf = mock.Mock()
+        driver = shotgun.driver.Driver({}, conf)
         result = driver.command(command)
         shotgun.driver.utils.execute.assert_called_with(command)
         self.assertEqual(result, out)
@@ -101,8 +123,9 @@ class TestDriver(base.BaseTestCase):
 
         command = "COMMAND"
 
+        conf = mock.Mock()
         driver = shotgun.driver.Driver(
-            {"host": {"address": "remote_host"}}, None)
+            {"host": {"address": "remote_host"}}, conf)
         result = driver.command(command)
 
         mstringio.assert_has_calls([
@@ -118,13 +141,14 @@ class TestDriver(base.BaseTestCase):
         mexecute.return_value = ("RETURN_CODE", "STDOUT", "STDERR")
         remote_path = "/remote_dir/remote_file"
         target_path = "/target_dir"
+        conf = mock.Mock()
 
         driver = shotgun.driver.Driver({
             "host": {
                 "address": "remote_host",
                 "ssh-key": "path_to_key",
             }
-        }, None)
+        }, conf)
         driver.get(remote_path, target_path)
         mexecute.assert_called_with('mkdir -p "{0}"'.format(target_path))
         mfabget.assert_called_with(remote_path, target_path)
@@ -133,12 +157,28 @@ class TestDriver(base.BaseTestCase):
             timeout=2, warn_only=True, abort_on_prompts=True)
 
         mexecute.reset_mock()
-        driver = shotgun.driver.Driver({}, None)
+        driver = shotgun.driver.Driver({}, conf)
         driver.get(remote_path, target_path)
         self.assertEqual(mexecute.mock_calls, [
             mock.call('mkdir -p "{0}"'.format(target_path)),
             mock.call('cp -r "{0}" "{1}"'.format(remote_path, target_path))])
 
+    def test_use_timeout_from_global_conf(self):
+        data = {}
+        conf = mock.Mock(spec=shotgun.config.Config, target="some_target")
+        cmd_driver = shotgun.driver.Driver(data, conf)
+        self.assertEqual(cmd_driver.timeout, conf.timeout)
+
+    def test_use_command_specific_timeout(self):
+        timeout = 1234
+        data = {
+            "timeout": timeout
+        }
+        conf = mock.Mock(spec=shotgun.config.Config, target="some_target")
+        cmd_driver = shotgun.driver.Driver(data, conf)
+        self.assertEqual(cmd_driver.timeout, timeout)
+        self.assertNotEqual(cmd_driver.timeout, conf.timeout)
+
 
 class TestFile(base.BaseTestCase):