Merge pull request #6865 from mvdbeek/share_key_function

Share ssh key generation across tests

test/base/ssh_util.py

@@ -0,0 +1,19 @@
+import collections
+import tempfile
+
+from Crypto.PublicKey import RSA
+
+
+def generate_ssh_keys():
+    """Returns a named tuple with private and public key and their paths."""
+    key = RSA.generate(2048)
+    private_key = key.export_key()
+    public_key = key.publickey().export_key(format='OpenSSH')
+    ssh_keys = collections.namedtuple('SSHKeys', 'private_key public_key private_key_file public_key_file')
+    with tempfile.NamedTemporaryFile(delete=False) as f:
+        f.write(private_key)
+        private_key_file = f.name
+    with tempfile.NamedTemporaryFile(delete=False) as f:
+        f.write(public_key)
+        public_key_file = f.name
+    return ssh_keys(private_key, public_key, private_key_file, public_key_file)

test/integration/test_cli_runners.py

@@ -6,30 +6,18 @@ import subprocess
 import tempfile
 import unittest
 
-from Crypto.PublicKey import RSA
-
 from base import integration_util  # noqa: I100,I202
 from base.populators import skip_without_tool
+from base.ssh_util import generate_ssh_keys
 from .test_job_environments import BaseJobEnvironmentIntegrationTestCase  # noqa: I201
 
 
-def generate_keys():
-    key = RSA.generate(2048)
-    return (key.export_key(), key.publickey().export_key(format='OpenSSH'))
-
-
 RemoteConnection = collections.namedtuple('remote_connection', ['hostname', 'username', 'password', 'port', 'private_key', 'public_key'])
 
 
 @integration_util.skip_unless_docker()
 def start_ssh_docker(container_name, jobs_directory, port=10022, image='agaveapi/slurm'):
-    private_key, public_key = generate_keys()
-    with tempfile.NamedTemporaryFile(delete=False) as f:
-        f.write(private_key)
-        private_key_file = f.name
-    with tempfile.NamedTemporaryFile(delete=False) as f:
-        f.write(public_key)
-        public_key_file = f.name
+    ssh_keys = generate_ssh_keys()
     START_SLURM_DOCKER = ['docker',
                           'run',
                           '-h',
@@ -43,12 +31,12 @@ def start_ssh_docker(container_name, jobs_directory, port=10022, image='agaveapi
                           '-v',
                           "{jobs_directory}:{jobs_directory}".format(jobs_directory=jobs_directory),
                           "-v",
-                          "{public_key_file}:/home/testuser/.ssh/authorized_keys".format(public_key_file=public_key_file),
+                          "{public_key_file}:/home/testuser/.ssh/authorized_keys".format(public_key_file=ssh_keys.public_key_file),
                           '--ulimit',
                           'nofile=2048:2048',
                           image]
     subprocess.check_call(START_SLURM_DOCKER)
-    return RemoteConnection('localhost', 'testuser', 'testuser', port, private_key_file, public_key_file)
+    return RemoteConnection('localhost', 'testuser', 'testuser', port, ssh_keys.private_key_file, ssh_keys.public_key_file)
 
 
 def stop_ssh_docker(container_name, remote_connection):
@@ -79,7 +67,7 @@ def cli_job_config(remote_connection, shell_plugin='ParamikoShell', job_plugin='
     job_conf_str = job_conf_template.substitute(shell_plugin=shell_plugin,
                                                 job_plugin=job_plugin,
                                                 **remote_connection._asdict())
-    with tempfile.NamedTemporaryFile(suffix="_slurm_integration_job_conf", delete=False) as job_conf:
+    with tempfile.NamedTemporaryFile(suffix="_slurm_integration_job_conf", mode="w", delete=False) as job_conf:
         job_conf.write(job_conf_str)
     return job_conf.name
 

test/unit/test_remote_shell.py

@@ -1,5 +1,4 @@
 import os
-import tempfile
 import unittest
 
 try:
@@ -7,7 +6,7 @@ try:
 except ImportError:
     raise unittest.SkipTest("Skipping tests that require mockssh")
 
-from Crypto.PublicKey import RSA
+from base.ssh_util import generate_ssh_keys
 
 from galaxy.jobs.runners.cli import CliInterface
 
@@ -16,19 +15,20 @@ class TestCliInterface(unittest.TestCase):
 
     @classmethod
     def tearDownClass(cls):
-        os.remove(cls.private_key)
+        os.remove(cls.ssh_keys.private_key_file)
+        os.remove(cls.ssh_keys.public_key_file)
 
     @classmethod
     def setUpClass(cls):
-        cls.private_key = make_private_key()
+        cls.ssh_keys = generate_ssh_keys()
         cls.username = 'testuser'
         cls.shell_params = {'username': cls.username,
-                            'private_key': cls.private_key,
+                            'private_key': cls.ssh_keys.private_key_file,
                             'hostname': 'localhost'}
         cls.cli_interface = CliInterface()
 
     def test_secure_shell_plugin_without_strict(self):
-        with mockssh.Server(users={self.username: self.private_key}) as server:
+        with mockssh.Server(users={self.username: self.ssh_keys.private_key_file}) as server:
             self.shell_params['port'] = server.port
             self.shell_params['plugin'] = 'SecureShell'
             self.shell_params['strict_host_key_checking'] = False
@@ -37,23 +37,16 @@ class TestCliInterface(unittest.TestCase):
         assert result.stdout.strip() == 'hello'
 
     def test_get_shell_plugin(self):
-        with mockssh.Server(users={self.username: self.private_key}) as server:
+        with mockssh.Server(users={self.username: self.ssh_keys.private_key_file}) as server:
             self.shell_params['port'] = server.port
             self.shell_params['plugin'] = 'ParamikoShell'
             self.shell = self.cli_interface.get_shell_plugin(self.shell_params)
         assert self.shell.username == self.username
 
     def test_paramiko_shell_plugin(self):
-        with mockssh.Server(users={self.username: self.private_key}) as server:
+        with mockssh.Server(users={self.username: self.ssh_keys.private_key_file}) as server:
             self.shell_params['port'] = server.port
             self.shell_params['plugin'] = 'ParamikoShell'
             self.shell = self.cli_interface.get_shell_plugin(self.shell_params)
             result = self.shell.execute(cmd='echo hello')
         assert result.stdout.strip() == 'hello'
-
-
-def make_private_key():
-    key = RSA.generate(1024)
-    private_fd, private_path = tempfile.mkstemp()
-    open(private_path, 'wb').write(key.exportKey('PEM'))
-    return private_path