diff --git a/shotgun/config.py b/shotgun/config.py index 2476bfe..e31a6b5 100644 --- a/shotgun/config.py +++ b/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) diff --git a/shotgun/driver.py b/shotgun/driver.py index e6dfdb9..d9c1886 100644 --- a/shotgun/driver.py +++ b/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"] diff --git a/shotgun/settings.py b/shotgun/settings.py index 8655c9f..36661e1 100644 --- a/shotgun/settings.py +++ b/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 diff --git a/shotgun/test/base.py b/shotgun/test/base.py index 1289e60..e0565ad 100644 --- a/shotgun/test/base.py +++ b/shotgun/test/base.py @@ -12,11 +12,7 @@ # License for the specific language governing permissions and limitations # under the License. -try: - from unittest.case import TestCase -except ImportError: - # Runing unit-tests in production environment - from unittest2.case import TestCase +from unittest2.case import TestCase class BaseTestCase(TestCase): diff --git a/shotgun/test/test_config.py b/shotgun/test/test_config.py index 2320cb4..448e593 100644 --- a/shotgun/test/test_config.py +++ b/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( - re.search( - ur"\/tmp\/sample\-[\d]{4}\-[\d]{2}\-[\d]{2}_" - "([\d]{2}\-){2}[\d]{2}", - conf.target - ) + self.assertRegex( + conf.target, + ur"\/tmp\/sample\-[\d]{4}\-[\d]{2}\-[\d]{2}_" + "([\d]{2}\-){2}[\d]{2}", ) + + @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) diff --git a/shotgun/test/test_driver.py b/shotgun/test/test_driver.py index 98d0385..85299c8 100644 --- a/shotgun/test/test_driver.py +++ b/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.driver -import shotgun.settings +import shotgun 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( - host_string="remote_host", timeout=2, command_timeout=10, - warn_only=True, key_filename=None, abort_on_prompts=True) + mfabset.assert_called_with( + host_string="remote_host", + 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):