From bd2c613b9c907d888b510a0a4c380e9fba55b69b Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 17 Sep 2017 21:09:03 +0530
Subject: [PATCH] Fix up a few more ssh_tests, stop sending 400 modules over
 network.

Defining functions in the current module was causing the entirety of
py.test and all dependencies to be sucked in.
---
 tests/data/plain_old_module.py |  9 ++++++
 tests/ssh_test.py              | 51 +++++++++++++---------------------
 2 files changed, 29 insertions(+), 31 deletions(-)

diff --git a/tests/data/plain_old_module.py b/tests/data/plain_old_module.py
index 5e4943d4..e09024a5 100755
--- a/tests/data/plain_old_module.py
+++ b/tests/data/plain_old_module.py
@@ -6,5 +6,14 @@ fiddlery.
 import math
 
 
+def get_sentinel_value():
+    # Some proof we're even talking to the mitogen-test Docker image
+    return file('/etc/sentinel').read()
+
+
+def add(x, y):
+    return x + y
+
+
 def pow(x, y):
     return x ** y
diff --git a/tests/ssh_test.py b/tests/ssh_test.py
index 035a315b..dabe0729 100644
--- a/tests/ssh_test.py
+++ b/tests/ssh_test.py
@@ -2,31 +2,23 @@
 import unittest
 
 import mitogen
-import mitogen.master
 import mitogen.ssh
 import mitogen.utils
 
 import testlib
+import plain_old_module
 
 
-def add(x, y):
-    return x + y
-
-
-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,
+class FakeSshTest(testlib.RouterMixin, unittest.TestCase):
+    def test_okay(self):
+        context = self.router.ssh(
                 hostname='hostname',
-                ssh_path=testlib.data_path('fakessh.py'))
-            context.call(mitogen.utils.log_to_file, '/tmp/log')
-            context.call(mitogen.utils.disable_site_packages)
-            self.assertEquals(3, context.call(add, 1, 2))
+                username='has-sudo',
+                ssh_path=testlib.data_path('fakessh.py'),
+        )
+        #context.call(mitogen.utils.log_to_file, '/tmp/log')
+        #context.call(mitogen.utils.disable_site_packages)
+        self.assertEquals(3, context.call(plain_old_module.add, 1, 2))
 
 
 class SshTest(testlib.DockerMixin, unittest.TestCase):
@@ -71,7 +63,7 @@ class SshTest(testlib.DockerMixin, unittest.TestCase):
         )
 
         sentinel = 'i-am-mitogen-test-docker-image\n'
-        assert sentinel == context.call(get_sentinel_value)
+        assert sentinel == context.call(plain_old_module.get_sentinel_value)
 
     def test_pubkey_required(self):
         try:
@@ -88,16 +80,13 @@ class SshTest(testlib.DockerMixin, unittest.TestCase):
         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
+        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 e[0] == self.stream_class.password_required_msg
+        sentinel = 'i-am-mitogen-test-docker-image\n'
+        assert sentinel == context.call(plain_old_module.get_sentinel_value)