From 9931a2e5184f86ede2acc100f8622120ca75f9f9 Mon Sep 17 00:00:00 2001
From: Thobias Salazar Trevisan <thobiast@gmail.com>
Date: Tue, 14 Jul 2020 13:35:02 -0300
Subject: [PATCH] Add context manager support to sshutils.SSH

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

Change-Id: If9cdf5c3c0e04041da89c1b60e97bfea2741b5cf
---
 rally/utils/sshutils.py           | 7 +++++++
 tests/unit/utils/test_sshutils.py | 9 +++++++++
 2 files changed, 16 insertions(+)

diff --git a/rally/utils/sshutils.py b/rally/utils/sshutils.py
index 33c043c16e..7b287eb6d9 100644
--- a/rally/utils/sshutils.py
+++ b/rally/utils/sshutils.py
@@ -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)
diff --git a/tests/unit/utils/test_sshutils.py b/tests/unit/utils/test_sshutils.py
index c9f5bdaeb0..555432041e 100644
--- a/tests/unit/utils/test_sshutils.py
+++ b/tests/unit/utils/test_sshutils.py
@@ -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()]