From cf506dfdf233752b82a9e5e73e92cc31d0d7f323 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 17 Sep 2017 18:11:48 +0530
Subject: [PATCH] First handful of functional tests for SSH against Docker.

---
 tests/ssh_test.py | 84 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 81 insertions(+), 3 deletions(-)

diff --git a/tests/ssh_test.py b/tests/ssh_test.py
index dac4bda0..035a315b 100644
--- a/tests/ssh_test.py
+++ b/tests/ssh_test.py
@@ -13,9 +13,13 @@ def add(x, y):
     return x + y
 
 
-class SshTest(unittest.TestCase):
-    def test_okay(self):
-        @mitogen.utils.run_with_broker
+def get_sentinel_value():
+    # Some proof we're even talking to the mitogen-test Docker image
+    return file('/etc/sentinel').read()
+
+
+class FakeSshTest(testlib.RouterMixin):
+    def test_okay(router, self):
         def test(broker):
             context = mitogen.ssh.connect(broker,
                 hostname='hostname',
@@ -23,3 +27,77 @@ class SshTest(unittest.TestCase):
             context.call(mitogen.utils.log_to_file, '/tmp/log')
             context.call(mitogen.utils.disable_site_packages)
             self.assertEquals(3, context.call(add, 1, 2))
+
+
+class SshTest(testlib.DockerMixin, unittest.TestCase):
+    stream_class = mitogen.ssh.Stream
+
+    def test_password_required(self):
+        try:
+            context = self.router.ssh(
+                hostname=self.dockerized_ssh.host,
+                port=self.dockerized_ssh.port,
+                check_host_keys=False,
+                username='has-sudo',
+            )
+            assert 0, 'exception not thrown'
+        except mitogen.ssh.PasswordError, e:
+            pass
+
+        assert e[0] == self.stream_class.password_required_msg
+
+    def test_password_incorrect(self):
+        try:
+            context = self.router.ssh(
+                hostname=self.dockerized_ssh.host,
+                port=self.dockerized_ssh.port,
+                check_host_keys=False,
+                username='has-sudo',
+                password='badpw',
+            )
+            assert 0, 'exception not thrown'
+        except mitogen.ssh.PasswordError, e:
+            pass
+
+        assert e[0] == self.stream_class.password_incorrect_msg
+
+    def test_password_specified(self):
+        context = self.router.ssh(
+            hostname=self.dockerized_ssh.host,
+            port=self.dockerized_ssh.port,
+            check_host_keys=False,
+            username='has-sudo',
+            password='y',
+        )
+
+        sentinel = 'i-am-mitogen-test-docker-image\n'
+        assert sentinel == context.call(get_sentinel_value)
+
+    def test_pubkey_required(self):
+        try:
+            context = self.router.ssh(
+                hostname=self.dockerized_ssh.host,
+                port=self.dockerized_ssh.port,
+                check_host_keys=False,
+                username='has-sudo-pubkey',
+            )
+            assert 0, 'exception not thrown'
+        except mitogen.ssh.PasswordError, e:
+            pass
+
+        assert e[0] == self.stream_class.password_required_msg
+
+    def test_pubkey_specified(self):
+        try:
+            context = self.router.ssh(
+                hostname=self.dockerized_ssh.host,
+                port=self.dockerized_ssh.port,
+                check_host_keys=False,
+                username='has-sudo-pubkey',
+                identity_file=testlib.data_path('docker/has-sudo-pubkey.key'),
+            )
+            assert 0, 'exception not thrown'
+        except mitogen.ssh.PasswordError, e:
+            pass
+
+        assert e[0] == self.stream_class.password_required_msg