[connection] Added upload method + mock-ssh-server tests

Author: nemesifier

Pipfile

@@ -14,6 +14,7 @@ mock = "*"
 pytest-django = ">=3.8.0,<4.0.0"
 pytest-asyncio = ">=0.10.0,<0.11.0"
 pytest-cov = ">=2.8.0,<2.9.0"
+mock-ssh-server = ">=0.8.0,<0.9.0"
 
 [scripts]
 lint = "python -m flake8"

openwisp_controller/connection/connectors/ssh.py

@@ -1,11 +1,12 @@
 import logging
 import socket
-from io import StringIO
+from io import BytesIO, StringIO
 
 import paramiko
 from django.utils.functional import cached_property
 from jsonschema import validate
 from jsonschema.exceptions import ValidationError as SchemaError
+from scp import SCPClient
 
 from .. import settings as app_settings
 
@@ -117,13 +118,11 @@ def exec_command(self, command, timeout=app_settings.SSH_COMMAND_TIMEOUT,
     def update_config(self):  # pragma: no cover
         raise NotImplementedError()
 
-    # TODO: this method is not used yet
-    # but will be necessary in the future to support other OSes
-    # def upload(self, fl, remote_path):
-    #     scp = SCPClient(self.shell.get_transport())
-    #     if not hasattr(fl, 'getvalue'):
-    #         fl_memory = BytesIO(fl.read())
-    #         fl.seek(0)
-    #         fl = fl_memory
-    #     scp.putfo(fl, remote_path)
-    #     scp.close()
+    def upload(self, fl, remote_path):
+        scp = SCPClient(self.shell.get_transport())
+        if not hasattr(fl, 'getvalue'):
+            fl_memory = BytesIO(fl.read())
+            fl.seek(0)
+            fl = fl_memory
+        scp.putfo(fl, remote_path)
+        scp.close()

openwisp_controller/connection/models.py

@@ -211,6 +211,7 @@ def connect(self):
         finally:
             self.last_attempt = timezone.now()
             self.save()
+        return self.is_working
 
     def disconnect(self):
         self.connector_instance.disconnect()

openwisp_controller/connection/tasks.py

@@ -13,7 +13,7 @@ def update_config(device_id):
     """
     # wait for the saving operations of this device to complete
     # (there may be multiple ones happening at the same time)
-    sleep(4)
+    sleep(2)
     # avoid repeating the operation multiple times
     device = Device.objects.select_related('config').get(pk=device_id)
     if device.config.status == 'applied':

openwisp_controller/connection/tests/base.py

@@ -14,6 +14,27 @@ class CreateConnectionsMixin(CreateConfigTemplateMixin, TestOrganizationMixin):
     device_model = Device
     config_model = Config
 
+    _TEST_RSA_PRIVATE_KEY_PATH = os.path.join(settings.BASE_DIR, 'test-key.rsa')
+    _TEST_RSA_PRIVATE_KEY_VALUE = None
+
+    class ssh_server:
+        host = '127.0.0.1'
+        port = 5555
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        with open(cls._TEST_RSA_PRIVATE_KEY_PATH, 'r') as f:
+            cls._TEST_RSA_PRIVATE_KEY_VALUE = f.read()
+
+    def _create_device(self, *args, **kwargs):
+        if 'last_ip' not in kwargs and 'management_ip' not in kwargs:
+            kwargs.update({
+                'last_ip': self.ssh_server.host,
+                'management_ip': self.ssh_server.host,
+            })
+        return super()._create_device(*args, **kwargs)
+
     def _create_credentials(self, **kwargs):
         opts = dict(name='Test credentials',
                     connector=app_settings.CONNECTORS[0][0],
@@ -31,7 +52,7 @@ def _create_credentials(self, **kwargs):
     def _create_credentials_with_key(self, username='root', port=22, **kwargs):
         opts = dict(name='Test SSH Key',
                     params={'username': username,
-                            'key': self._SSH_PRIVATE_KEY,
+                            'key': self._TEST_RSA_PRIVATE_KEY_VALUE,
                             'port': port})
         opts.update(kwargs)
         return self._create_credentials(**opts)
@@ -53,18 +74,3 @@ def _create_device_connection(self, **kwargs):
         dc.full_clean()
         dc.save()
         return dc
-
-
-class SshMixin(object):
-    _TEST_RSA_KEY_PATH = os.path.join(settings.BASE_DIR, 'test-key.rsa')
-    _SSH_PRIVATE_KEY = None
-
-    class ssh_server:
-        host = '127.0.0.1'
-        port = 5555
-
-    @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
-        with open(cls._TEST_RSA_KEY_PATH, 'r') as f:
-            cls._SSH_PRIVATE_KEY = f.read()

openwisp_controller/connection/tests/test_admin.py

@@ -5,11 +5,10 @@
 from ...config.tests.test_admin import TestAdmin as TestConfigAdmin
 from ...tests.utils import TestAdminMixin
 from ..models import Credentials, DeviceConnection
-from .base import CreateConnectionsMixin, SshMixin
+from .base import CreateConnectionsMixin
 
 
-class TestAdmin(TestAdminMixin, CreateConnectionsMixin,
-                SshMixin, TestCase):
+class TestAdmin(TestAdminMixin, CreateConnectionsMixin, TestCase):
     template_model = Template
     credentials_model = Credentials
     connection_model = DeviceConnection

openwisp_controller/connection/tests/test_models.py

@@ -9,20 +9,12 @@
 
 from .. import settings as app_settings
 from ..models import Credentials
-from .base import CreateConnectionsMixin, SshMixin
+from .base import CreateConnectionsMixin
 
 
-class TestModels(SshMixin, CreateConnectionsMixin, TestCase):
+class TestModels(CreateConnectionsMixin, TestCase):
     _connect_path = 'paramiko.SSHClient.connect'
 
-    def _create_device(self, *args, **kwargs):
-        if 'last_ip' not in kwargs and 'management_ip' not in kwargs:
-            kwargs.update({
-                'last_ip': self.ssh_server.host,
-                'management_ip': self.ssh_server.host,
-            })
-        return super()._create_device(*args, **kwargs)
-
     def test_connection_str(self):
         c = Credentials(name='Dev Key', connector=app_settings.CONNECTORS[0][0])
         self.assertIn(c.name, str(c))

openwisp_controller/connection/tests/test_ssh.py

@@ -0,0 +1,54 @@
+import io
+import os
+from contextlib import redirect_stdout
+
+import mock
+from django.conf import settings
+from django.test import TestCase
+from mockssh import Server
+
+from .base import CreateConnectionsMixin
+
+
+class TestSsh(CreateConnectionsMixin, TestCase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.mock_ssh_server = Server({'root': cls._TEST_RSA_PRIVATE_KEY_PATH}).__enter__()
+        cls.ssh_server.port = cls.mock_ssh_server.port
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.mock_ssh_server.__exit__()
+
+    def test_connection_connect(self):
+        ckey = self._create_credentials_with_key(port=self.ssh_server.port)
+        dc = self._create_device_connection(credentials=ckey)
+        dc.connector_instance.connect()
+        stdout = io.StringIO()
+        with redirect_stdout(stdout):
+            dc.connector_instance.exec_command('echo test')
+        output = stdout.getvalue()
+        self.assertIn('$:> echo test\ntest', output)
+
+    def test_connection_failed_command(self):
+        ckey = self._create_credentials_with_key(port=self.ssh_server.port)
+        dc = self._create_device_connection(credentials=ckey)
+        dc.connector_instance.connect()
+        stdout = io.StringIO()
+        with redirect_stdout(stdout):
+            with self.assertRaises(Exception):
+                dc.connector_instance.exec_command('wrongcommand')
+        output = stdout.getvalue()
+        self.assertIn('/bin/sh: 1: wrongcommand: not found', output)
+        self.assertIn('# Previus command failed, aborting...', output)
+
+    @mock.patch('scp.SCPClient.putfo')
+    def test_connection_upload(self, putfo_mocked):
+        ckey = self._create_credentials_with_key(port=self.ssh_server.port)
+        dc = self._create_device_connection(credentials=ckey)
+        dc.connector_instance.connect()
+        # needs a binary file to test all lines
+        fl = open(os.path.join(settings.BASE_DIR, '../media/floorplan.jpg'), 'rb')
+        dc.connector_instance.upload(fl, '/tmp/test')
+        putfo_mocked.assert_called_once()