Turn shells into context managers

spur/local.py

@@ -9,6 +9,12 @@
 
 
 class LocalShell(object):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
     def upload_dir(self, source, dest, ignore=None):
         shutil.copytree(source, dest, ignore=shutil.ignore_patterns(*ignore))
 

spur/ssh.py

@@ -32,6 +32,15 @@ def __init__(self, hostname, username, password=None, port=22, private_key_file=
         self._private_key_file = private_key_file
         self._client = None
         self._connect_timeout = connect_timeout if not None else _ONE_MINUTE
+        self._closed = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self._closed = True
+        if self._client is not None:
+            self._client.close()
 
     def run(self, *args, **kwargs):
         return self.spawn(*args, **kwargs).wait_for_result()
@@ -136,6 +145,8 @@ def _get_ssh_transport(self):
 
     def _connect_ssh(self):
         if self._client is None:
+            if self._closed:
+                raise RuntimeError("Shell is closed")
             client = paramiko.SSHClient()
             client.load_system_host_keys()
             client.set_missing_host_key_policy(paramiko.WarningPolicy())

tests/__init__.py

@@ -0,0 +1 @@
+

tests/ssh_tests.py

@@ -2,6 +2,7 @@
 
 import spur
 import spur.ssh
+from .testing import create_ssh_shell
 
 
 @istest
@@ -17,3 +18,10 @@ def try_connection():
 
     assert_raises(spur.ssh.ConnectionError, try_connection)
 
+
+@istest
+def trying_to_use_ssh_shell_after_exit_results_in_error():
+    with create_ssh_shell() as shell:
+        pass
+
+    assert_raises(Exception, lambda: shell.run(["true"]))

tests/testing.py

@@ -0,0 +1,14 @@
+import os
+
+import spur
+
+
+def create_ssh_shell():
+    port_var = os.environ.get("TEST_SSH_PORT")
+    port = int(port_var) if port_var is not None else None
+    return spur.SshShell(
+        hostname=os.environ.get("TEST_SSH_HOSTNAME", "127.0.0.1"),
+        username=os.environ["TEST_SSH_USERNAME"],
+        password=os.environ["TEST_SSH_PASSWORD"],
+        port=port
+    )

tests/tests.py

@@ -8,29 +8,22 @@
 from nose.tools import istest, assert_equal, assert_raises, assert_true
 
 import spur
+from .testing import create_ssh_shell
+
 
 def test(func):
     @functools.wraps(func)
     def run_test():
         for shell in _create_shells():
-            yield func, shell
+            with shell:
+                yield func, shell
 
     def _create_shells():
         return [
             spur.LocalShell(),
-            _create_ssh_shell()
+            create_ssh_shell()
         ]
 
-    def _create_ssh_shell():
-        port_var = os.environ.get("TEST_SSH_PORT")
-        port = int(port_var) if port_var is not None else None
-        return spur.SshShell(
-            hostname=os.environ.get("TEST_SSH_HOSTNAME", "127.0.0.1"),
-            username=os.environ["TEST_SSH_USERNAME"],
-            password=os.environ["TEST_SSH_PASSWORD"],
-            port=port
-        )
-
     return istest(run_test)
 
 @test