From 4166448427f459dec59c1121e9f3ce49dd7b8317 Mon Sep 17 00:00:00 2001 From: Alexey Stepanov Date: Thu, 18 Aug 2016 08:29:49 +0300 Subject: [PATCH] Rework get_sudo get_sudo is not used in fuel-qa starting from 6.1 up to master, so change is safe get_sudo context manager -> mangle add sudo method with ling to object Add parameter enforce, this allow to enforce enable or disable sudo (and do not change by default) Cover by unit tests Change-Id: I74105cf251872f394a23867b0aa2c37906e7abf0 --- devops/helpers/ssh_client.py | 41 ++++++++-- devops/tests/helpers/test_ssh_client.py | 102 +++++++++++++++++++++++- 2 files changed, 135 insertions(+), 8 deletions(-) diff --git a/devops/helpers/ssh_client.py b/devops/helpers/ssh_client.py index 1386ee80..5002e953 100644 --- a/devops/helpers/ssh_client.py +++ b/devops/helpers/ssh_client.py @@ -253,6 +253,7 @@ class _MemorizedSSH(type): logger.debug('Closing {} as unused'.format(cls.__cache[key])) cls.__cache[key].close() del cls.__cache[key] + # noinspection PyArgumentList return super( _MemorizedSSH, cls).__call__( host=host, port=port, @@ -305,21 +306,41 @@ class SSHClient(six.with_metaclass(_MemorizedSSH, object)): class get_sudo(object): """Context manager for call commands with sudo""" + def __init__(self, ssh): + warn( + 'SSHClient.get_sudo(SSHClient()) is deprecated in favor of ' + 'SSHClient().sudo(enforce=...) , which is much more powerful.') self.ssh = ssh self.__sudo_status = False - def __enter__(self, enable_sudo=True): - """Context manager for handling sudo mode - - :type enable_sudo: bool - """ + def __enter__(self): self.__sudo_status = self.ssh.sudo_mode - self.ssh.sudo_mode = enable_sudo + self.ssh.sudo_mode = True def __exit__(self, exc_type, exc_val, exc_tb): self.ssh.sudo_mode = self.__sudo_status + class __get_sudo(object): + """Context manager for call commands with sudo""" + def __init__(self, ssh, enforce=None): + """Context manager for call commands with sudo + + :type ssh: SSHClient + :type enforce: bool + """ + self.__ssh = ssh + self.__sudo_status = ssh.sudo_mode + self.__enforce = enforce + + def __enter__(self): + self.__sudo_status = self.__ssh.sudo_mode + if self.__enforce is not None: + self.__ssh.sudo_mode = self.__enforce + + def __exit__(self, exc_type, exc_val, exc_tb): + self.__ssh.sudo_mode = self.__sudo_status + def __hash__(self): return hash(( self.__class__, @@ -551,6 +572,14 @@ class SSHClient(six.with_metaclass(_MemorizedSSH, object)): self.__connect() + def sudo(self, enforce=None): + """Call contextmanager for sudo mode change + + :type enforce: bool + :param enforce: Enforce sudo enabled or disabled. By default: None + """ + return self.__get_sudo(ssh=self, enforce=enforce) + def check_call( self, command, verbose=False, timeout=None, diff --git a/devops/tests/helpers/test_ssh_client.py b/devops/tests/helpers/test_ssh_client.py index 7e63b83d..f4b6b627 100644 --- a/devops/tests/helpers/test_ssh_client.py +++ b/devops/tests/helpers/test_ssh_client.py @@ -1008,7 +1008,7 @@ class TestExecute(TestCase): logger.mock_calls ) - def test_execute_async_with_sudo(self, client, policy, logger): + def test_execute_async_with_sudo_enforce(self, client, policy, logger): chan = mock.Mock() open_session = mock.Mock(return_value=chan) transport = mock.Mock() @@ -1020,7 +1020,7 @@ class TestExecute(TestCase): ssh = self.get_ssh() self.assertFalse(ssh.sudo_mode) - with SSHClient.get_sudo(ssh): + with SSHClient.sudo(ssh, enforce=True): self.assertTrue(ssh.sudo_mode) # noinspection PyTypeChecker result = ssh.execute_async(command=command) @@ -1044,6 +1044,70 @@ class TestExecute(TestCase): logger.mock_calls ) + def test_execute_async_with_no_sudo_enforce(self, client, policy, logger): + chan = mock.Mock() + open_session = mock.Mock(return_value=chan) + transport = mock.Mock() + transport.attach_mock(open_session, 'open_session') + get_transport = mock.Mock(return_value=transport) + _ssh = mock.Mock() + _ssh.attach_mock(get_transport, 'get_transport') + client.return_value = _ssh + + ssh = self.get_ssh() + ssh.sudo_mode = True + + with ssh.sudo(enforce=False): + # noinspection PyTypeChecker + result = ssh.execute_async(command=command) + get_transport.assert_called_once() + open_session.assert_called_once() + + self.assertIn(chan, result) + chan.assert_has_calls(( + mock.call.makefile('wb'), + mock.call.makefile('rb'), + mock.call.makefile_stderr('rb'), + mock.call.exec_command('{}\n'.format(command)) + )) + self.assertIn( + mock.call.debug( + "Executing command: '{}'".format(command.rstrip())), + logger.mock_calls + ) + + def test_execute_async_with_none_enforce(self, client, policy, logger): + chan = mock.Mock() + open_session = mock.Mock(return_value=chan) + transport = mock.Mock() + transport.attach_mock(open_session, 'open_session') + get_transport = mock.Mock(return_value=transport) + _ssh = mock.Mock() + _ssh.attach_mock(get_transport, 'get_transport') + client.return_value = _ssh + + ssh = self.get_ssh() + ssh.sudo_mode = False + + with ssh.sudo(): + # noinspection PyTypeChecker + result = ssh.execute_async(command=command) + get_transport.assert_called_once() + open_session.assert_called_once() + + self.assertIn(chan, result) + chan.assert_has_calls(( + mock.call.makefile('wb'), + mock.call.makefile('rb'), + mock.call.makefile_stderr('rb'), + mock.call.exec_command('{}\n'.format(command)) + )) + self.assertIn( + mock.call.debug( + "Executing command: '{}'".format(command.rstrip())), + logger.mock_calls + ) + @mock.patch('devops.helpers.ssh_client.SSHAuth.enter_password') def test_execute_async_sudo_password( self, enter_password, client, policy, logger): @@ -1315,6 +1379,40 @@ class TestExecute(TestCase): ssh.check_call(command=command, verbose=verbose, timeout=None) execute.assert_called_once_with(command, verbose, None) + @mock.patch( + 'devops.helpers.ssh_client.SSHClient.execute') + def test_check_call_expected(self, execute, client, policy, logger): + exit_code = 0 + return_value = { + 'stderr_str': '0\n1', + 'stdout_str': '2\n3', + 'exit_code': exit_code, + 'stderr': [b' \n', b'0\n', b'1\n', b' \n'], + 'stdout': [b' \n', b'2\n', b'3\n', b' \n']} + execute.return_value = return_value + + verbose = False + + ssh = self.get_ssh() + + # noinspection PyTypeChecker + result = ssh.check_call( + command=command, verbose=verbose, timeout=None, expected=[0, 75]) + execute.assert_called_once_with(command, verbose, None) + self.assertEqual(result, return_value) + + exit_code = 1 + return_value['exit_code'] = exit_code + execute.reset_mock() + execute.return_value = return_value + with self.assertRaises(DevopsCalledProcessError): + # noinspection PyTypeChecker + ssh.check_call( + command=command, verbose=verbose, timeout=None, + expected=[0, 75] + ) + execute.assert_called_once_with(command, verbose, None) + @mock.patch( 'devops.helpers.ssh_client.SSHClient.check_call') def test_check_stderr(self, check_call, client, policy, logger):