Convert SSHClient into a context manager

1 files changed, 24 insertions, 0 deletions

diff --git a/tests/test_client.py b/tests/test_client.py
index 7c094628..b3635272 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -20,6 +20,8 @@
 Some unit tests for SSHClient.
 """
 
+from __future__ import with_statement
+
 import socket
 from tempfile import mkstemp
 import threading
@@ -293,3 +295,25 @@ class SSHClientTest (unittest.TestCase):
         gc.collect()
 
         self.assertTrue(p() is None)
+
+    def test_6_client_can_be_used_as_context_manager(self):
+        """
+        verify that an SSHClient can be used a context manager
+        """
+        threading.Thread(target=self._run).start()
+        host_key = paramiko.RSAKey.from_private_key_file(test_path('test_rsa.key'))
+        public_host_key = paramiko.RSAKey(data=host_key.asbytes())
+
+        with paramiko.SSHClient() as tc:
+            self.tc = tc
+            self.tc.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            self.assertEquals(0, len(self.tc.get_host_keys()))
+            self.tc.connect(self.addr, self.port, username='slowdive', password='pygmalion')
+
+            self.event.wait(1.0)
+            self.assertTrue(self.event.isSet())
+            self.assertTrue(self.ts.is_active())
+
+            self.assertTrue(self.tc._transport is not None)
+
+        self.assertTrue(self.tc._transport is None)