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:
parent
695f5ecbab
commit
9931a2e518
@ -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)
|
||||||
|
@ -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()]
|
||||||
|
Loading…
Reference in New Issue
Block a user