openstack/rally
Code Issues Proposed changes

Add context manager support to sshutils.SSH

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
rally/utils/sshutils.py
View File

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

9
tests/unit/utils/test_sshutils.py
View File

@ -119,6 +119,15 @@ class SSHTestCase(test.TestCase):
m_client.close.assert_called_once_with()
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")
def test_execute(self, mock_string_io):
mock_string_io.side_effect = stdio = [mock.Mock(), mock.Mock()]