Merge "Add context manager support to sshutils.SSH"

rally/utils/sshutils.py

@@ -95,6 +95,13 @@ def __init__(self, user, host, port=22, pkey=None,
         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)

tests/unit/utils/test_sshutils.py

@@ -119,6 +119,15 @@ def test_close(self):
         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()]