Add test on notification management, fix improper message when user upload the notification file

mod_config/controllers.py

@@ -143,7 +143,7 @@ def notifications():
                         # Delayed insert, show message for delay
                         form.errors['file'] = [
                             'The file was submitted, but dependencies need '
-                            'to be installed. It will be listed as soon as '
+                            'to be installed. It will be listed after refreshing the page if'
                             'said dependencies are installed.'
                         ]
                 except NotificationLoader.NotificationLoaderException as e:

mod_config/forms.py

@@ -44,6 +44,7 @@ def validate_file(form, field):
 
 def simple_notification_file_validation(check_notification=True):
     def validate_file(form, field):
+        field.data.filename = os.path.basename(field.data.filename)
         file_type = is_python_or_container(field.data.filename)
         if file_type == FileType.PYTHONFILE:
             # Name cannot be one of the files we already have

tests/testAppBase.py

@@ -70,15 +70,17 @@ def create_app(self):
 
     def create_admin(self):
         # test if there is admin existed
-        name, password, email = "admin", "adminpwd", "admin@email.com"
-        db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
-        role = Role(name=name)
-        db.add(role)
-        db.commit()
-        admin_user = User(role_id=role.id, name=name, password=password, email=email)
-        db.add(admin_user)
-        db.commit()
-        db.remove()
+        try:
+            name, password, email = "admin", "adminpwd", "admin@email.com"
+            db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
+            role = Role(name=name)
+            db.add(role)
+            db.commit()
+            admin_user = User(role_id=role.id, name=name, password=password, email=email)
+            db.add(admin_user)
+            db.commit()
+        finally:
+            db.remove()
         return name, password, email
 
     def setUp(self):
@@ -89,10 +91,12 @@ def setUp(self):
 
     def tearDown(self):
         from database import Base
-        db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
-        db_engine = create_engine(self.app.config['DATABASE_URI'], convert_unicode=True)
-        Base.metadata.drop_all(bind=db_engine)
-        db.remove()
+        try:
+            db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
+            db_engine = create_engine(self.app.config['DATABASE_URI'], convert_unicode=True)
+            Base.metadata.drop_all(bind=db_engine)
+        finally:
+            db.remove()
 
 
 class TestApp(TestAppBase):

tests/testFiles/ModifiedTelegramNotification/TelegramNotification.py

@@ -0,0 +1,39 @@
+from pipot.notifications.INotification import INotification
+
+
+class TelegramNotification(INotification):
+    def __init__(self, config):
+        super(TelegramNotification, self).__init__(config)
+
+    def process(self, message):
+        for chat_id in self.config['chat_ids']:
+            self.bot.send_message(chat_id, message)
+
+    @staticmethod
+    def get_pip_dependencies():
+        return ['python-telegram-bot']
+
+    @staticmethod
+    def after_install_hook():
+        return True
+
+    @staticmethod
+    def get_apt_dependencies():
+        return []
+
+    def requires_extra_config(self):
+        return True
+
+    def is_valid_extra_config(self, config):
+        return 'token' in config and 'chat_ids' in config
+
+    @staticmethod
+    def get_extra_config_sample():
+        return {
+            'token': '123456:ABC',
+            'chat_ids': [123456, 1234567]
+        }
+
+    @staticmethod
+    def modified_function_for_test():
+        pass

tests/testFiles/TelegramNotification/TelegramNotification.py

@@ -0,0 +1,35 @@
+from pipot.notifications.INotification import INotification
+
+
+class TelegramNotification(INotification):
+    def __init__(self, config):
+        super(TelegramNotification, self).__init__(config)
+
+    def process(self, message):
+        for chat_id in self.config['chat_ids']:
+            self.bot.send_message(chat_id, message)
+
+    @staticmethod
+    def get_pip_dependencies():
+        return ['python-telegram-bot']
+
+    @staticmethod
+    def after_install_hook():
+        return True
+
+    @staticmethod
+    def get_apt_dependencies():
+        return []
+
+    def requires_extra_config(self):
+        return True
+
+    def is_valid_extra_config(self, config):
+        return 'token' in config and 'chat_ids' in config
+
+    @staticmethod
+    def get_extra_config_sample():
+        return {
+            'token': '123456:ABC',
+            'chat_ids': [123456, 1234567]
+        }

tests/testNotificationManagement.py

@@ -0,0 +1,155 @@
+import os
+import sys
+import mock
+import unittest
+import json
+import codecs
+import filecmp
+from mock import patch
+from functools import wraps
+from werkzeug.datastructures import FileStorage
+
+from flask import request, jsonify
+
+import tests.authMock
+from database import create_session
+from mod_config.models import Notification
+from tests.testAppBase import TestAppBase
+
+test_dir = os.path.dirname(os.path.abspath(__file__))
+notification_dir = os.path.join(test_dir, '../pipot/notifications/')
+temp_dir = os.path.join(test_dir, 'temp')
+if not os.path.exists(temp_dir):
+    os.makedirs(temp_dir)
+
+
+def _install_notification_service(cls, update=True, description=""):
+    from mod_config.controllers import g
+    if not update:
+        instance = cls(getattr(cls, 'get_extra_config_sample')())
+        notification = Notification(
+            instance.__class__.__name__, description)
+        g.db.add(notification)
+        g.db.commit()
+    return True
+
+
+class TestNotificationManagement(TestAppBase):
+
+    def setUp(self):
+        super(TestNotificationManagement, self).setUp()
+
+    def tearDown(self):
+        super(TestNotificationManagement, self).tearDown()
+
+    @patch("mod_config.controllers._install_notification_service", side_effect=_install_notification_service)
+    def add_notification(self, notification_name, notification_file_name, install_mock):
+        # upload the notification file
+        notification_file = codecs.open(os.path.join(test_dir, 'testFiles', notification_name,
+                                                     notification_file_name), 'rb')
+        # notification_file = FileStorage(notification_file)
+        with self.app.test_client() as client:
+            data = dict(
+                file=notification_file,
+                description='test'
+            )
+            response = client.post('/notifications', data=data, follow_redirects=False)
+            self.assertEqual(response.status_code, 200)
+        install_mock.assert_called_once()
+        # check backup notification file is removed under temp_path
+        self.assertFalse(os.path.isfile(os.path.join(notification_dir, 'temp', notification_file_name)))
+        # check notification file and folder is created under final_path
+        self.assertTrue(os.path.isfile(os.path.join(notification_dir, notification_name + '.py')))
+        # check database
+        try:
+            db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
+            notification_row = db.query(Notification.id, Notification.name).first()
+            notification_id = notification_row.id
+            name = notification_row.name
+        finally:
+            db.remove()
+        self.assertEqual(notification_name, name)
+        return notification_id
+
+    def remove_notification(self, notification_id, notification_name):
+        # delete notification file
+        with self.app.test_client() as client:
+            data = dict(
+                id=notification_id
+            )
+            response = client.post('/notifications/delete', data=data, follow_redirects=False)
+            self.assertEqual(response.status_code, 200)
+            self.assertEqual(response.get_json()['status'], 'success')
+        # check notification file and folder is removed unser final_path
+        self.assertFalse(os.path.isfile(os.path.join(notification_dir, notification_name + '.py')))
+        # check database
+        try:
+            db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
+            notification_row = db.query(Notification.id).first()
+            result = (notification_row is None)
+        finally:
+            db.remove()
+        self.assertTrue(result)
+
+    @patch("mod_config.controllers._install_notification_service", side_effect=_install_notification_service)
+    def update_notification(self, notification_id, notification_name, notification_file_name, install_mock):
+        notification_file = codecs.open(os.path.join(test_dir, 'testFiles', notification_file_name), 'rb')
+        # update notification file
+        with self.app.test_client() as client:
+            data = dict(
+                notificationUpdate_id=notification_id,
+                notificationUpdate_file=notification_file,
+            )
+            response = client.post('/notifications/update', data=data, follow_redirects=False)
+            self.assertEqual(response.status_code, 200)
+        install_mock.assert_called_once()
+        # check backup notification file is removed under temp_path
+        self.assertFalse(os.path.isfile(os.path.join(notification_dir, 'temp', notification_name + '.py')))
+        # check notification file and folder still exist under final path
+        self.assertTrue(os.path.isfile(os.path.join(notification_dir, notification_name + '.py')))
+        return response
+
+    def test_add_and_delete_notification_file(self):
+        notification_name = 'TelegramNotification'
+        notification_file_name = notification_name + '.py'
+        notification_id = self.add_notification(notification_name, notification_file_name)
+        self.remove_notification(notification_id, notification_name)
+
+    def test_update_with_valid_notification_file(self):
+        notification_name = 'TelegramNotification'
+        notification_file_name = notification_name + '.py'
+        # add a new discription column
+        modified_notification_file_namae = 'ModifiedTelegramNotification/TelegramNotification.py'
+        notification_id = self.add_notification(notification_name, notification_file_name)
+        response = self.update_notification(notification_id, notification_name, modified_notification_file_namae)
+        self.assertEqual(response.get_json()['status'], 'success')
+        # check file content
+        self.assertTrue(filecmp.cmp(os.path.join(notification_dir, notification_file_name),
+                                    os.path.join(test_dir, 'testFiles', modified_notification_file_namae)))
+        # check on database
+        try:
+            db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
+            notification_row = db.query(Notification.id, Notification.name).first()
+            notification_id = notification_row.id
+            name = notification_row.name
+        finally:
+            db.remove()
+        self.assertEqual(notification_name, name)
+        self.remove_notification(notification_id, notification_name)
+
+    # def test_update_with_invalid_notification_file(self):
+    #     notification_name = 'TelnetService'
+    #     notification_file_name = notification_name + '.py'
+    #     # try to update an invalid notification file
+    #     modified_notification_file_name = 'EmptyTelnetService/TelnetService.py'
+    #     notification_id = self.add_notification(notification_name, notification_file_name)
+    #     response = self.update_notification(notification_id, notification_name, modified_notification_file_name)
+    #     # check the notification file doesn't change
+    #     self.assertTrue(filecmp.cmp(os.path.join(notification_dir, notification_name, notification_name + '.py'),
+    #                     os.path.join(test_dir, 'testFiles', notification_file_name)))
+    #     self.assertEqual(response.get_json()['status'], 'error')
+    #     self.remove_notification(notification_id, notification_name)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/testServiceManagement.py

@@ -65,7 +65,7 @@ def add_service(self, service_name, service_file_name):
         self.assertEqual(service_name, name)
         return service_id
 
-    def remove_service(self, service_id, service_name, service_file_name):
+    def remove_service(self, service_id, service_name):
         # delete service file
         with self.app.test_client() as client:
             data = dict(
@@ -75,7 +75,7 @@ def remove_service(self, service_id, service_name, service_file_name):
             self.assertEqual(response.status_code, 200)
             self.assertEqual(response.get_json()['status'], 'success')
         # check service file and folder is removed unser final_path
-        self.assertFalse(os.path.isfile(os.path.join(service_dir, service_name, service_file_name)))
+        self.assertFalse(os.path.isfile(os.path.join(service_dir, service_name, service_name + '.py')))
         # check models.txt is updated
         self.assertEqual([], ServiceModelsManager.get_models())
         # check database
@@ -101,32 +101,33 @@ def update_service(self, service_id, service_name, service_file_name):
             )
             response = client.post('/services/update', data=data, follow_redirects=False)
             self.assertEqual(response.status_code, 200)
-            self.assertEqual(response.get_json()['status'], 'success')
         # check backup service file is removed under temp_path
         self.assertFalse(os.path.isfile(os.path.join(service_dir, 'temp', service_file_name)))
         # check service file and folder still exist under final path
         self.assertTrue(os.path.isdir(os.path.join(service_dir, service_name)))
         self.assertTrue(os.path.isfile(os.path.join(service_dir, service_name, service_name + '.py')))
+        return response
 
     def test_add_and_delete_service_file(self):
         service_name = 'TelnetService'
         service_file_name = service_name + '.py'
         service_id = self.add_service(service_name, service_file_name)
-        self.remove_service(service_id, service_name, service_file_name)
+        self.remove_service(service_id, service_name)
 
     def test_add_and_delete_service_container(self):
         service_name = 'TelnetService'
         service_file_name = service_name + '.zip'
         service_id = self.add_service(service_name, service_file_name)
-        self.remove_service(service_id, service_name, service_file_name)
+        self.remove_service(service_id, service_name)
 
     def test_update_with_valid_service_file(self):
         service_name = 'TelnetService'
         service_file_name = service_name + '.py'
         # add a new discription column
-        modified_service_file_namae = 'ModifiedService/TelnetService.py'
+        modified_service_file_namae = 'ModifiedTelnetService/TelnetService.py'
         service_id = self.add_service(service_name, service_file_name)
-        self.update_service(service_id, service_name, modified_service_file_namae)
+        response = self.update_service(service_id, service_name, modified_service_file_namae)
+        self.assertEqual(response.get_json()['status'], 'success')
         # check on metadata
         from database import Base
         has_table = False
@@ -136,7 +137,20 @@ def test_update_with_valid_service_file(self):
                 self.assertTrue('description' in table.columns.keys())
                 break
         self.assertTrue(has_table)
-        self.remove_service(service_id, service_name, service_file_name)
+        self.remove_service(service_id, service_name)
+
+    def test_update_with_invalid_service_file(self):
+        service_name = 'TelnetService'
+        service_file_name = service_name + '.py'
+        # try to update an invalid service file
+        modified_service_file_name = 'EmptyTelnetService/TelnetService.py'
+        service_id = self.add_service(service_name, service_file_name)
+        response = self.update_service(service_id, service_name, modified_service_file_name)
+        # check the service file doesn't change
+        self.assertTrue(filecmp.cmp(os.path.join(service_dir, service_name, service_name + '.py'),
+                        os.path.join(test_dir, 'testFiles', service_file_name)))
+        self.assertEqual(response.get_json()['status'], 'error')
+        self.remove_service(service_id, service_name)
 
 
 if __name__ == '__main__':