Add context manager support to sshutils.SSH

Enable the use of context managers for sshutils.SSH class.

Change-Id: If9cdf5c3c0e04041da89c1b60e97bfea2741b5cf
This commit is contained in:
Thobias Salazar Trevisan 2020-07-14 13:35:02 -03:00
parent 695f5ecbab
commit 9931a2e518
2 changed files with 16 additions and 0 deletions

View File

@ -95,6 +95,13 @@ class SSH(object):
self.key_filename = key_filename self.key_filename = key_filename
self._client = False self._client = False
def __enter__(self):
return self
def __exit__(self, *args):
if self._client:
self.close()
def _get_pkey(self, key): def _get_pkey(self, key):
if isinstance(key, str): if isinstance(key, str):
key = io.StringIO(key) key = io.StringIO(key)

View File

@ -119,6 +119,15 @@ class SSHTestCase(test.TestCase):
m_client.close.assert_called_once_with() m_client.close.assert_called_once_with()
self.assertFalse(self.ssh._client) self.assertFalse(self.ssh._client)
def test_close_context_manager_enter(self):
self.assertEqual(self.ssh, self.ssh.__enter__())
def test_close_context_manager_exit(self):
with mock.patch.object(self.ssh, "_client") as m_client:
self.ssh.__exit__()
m_client.close.assert_called_once_with()
self.assertFalse(self.ssh._client)
@mock.patch("rally.utils.sshutils.io.StringIO") @mock.patch("rally.utils.sshutils.io.StringIO")
def test_execute(self, mock_string_io): def test_execute(self, mock_string_io):
mock_string_io.side_effect = stdio = [mock.Mock(), mock.Mock()] mock_string_io.side_effect = stdio = [mock.Mock(), mock.Mock()]