Avoid imp usage in monitoring and worker tests

tests/test_mediaworker.py

@@ -4,43 +4,25 @@
 Criticality: High
 Checks that MediaWorker can be reached using SSH and that it can reach the tasks server
 '''
-import imp
-import os
+from pathlib import Path
 import re
 import subprocess
 import sys
 
-GREEN = '\033[92m'
-RED = '\033[91m'
-DEF = '\033[0m'
+sys.path.append(str(Path(__file__).parents[1].resolve()))
 
-
-try:
-    import mediaserver
-except ImportError:
-    print('MediaServer is not installed, skipping test.')
-    sys.exit(2)
-else:
-    print('MediaServer version: %s.' % mediaserver.__version__)
-
-
-os.chdir(os.path.dirname(__file__))
-if not os.path.isfile('../utils.py'):
-    print('conf.sh not found.')
-    sys.exit(1)
-
-es_utils = imp.load_source('es_utils', '../utils.py')
-conf = es_utils.load_conf()
+# pylint: disable=wrong-import-position
+from envsetup import utils as u  # noqa: E402
 
 
 def check_ssh(ip):
     cmd = 'ssh -o StrictHostKeyChecking=no -o PasswordAuthentication=no root@%s ls /tmp' % ip
-    print('Connecting to MediaWorker:\n%s' % cmd)
+    u.log('Connecting to MediaWorker:\n%s' % cmd)
     try:
         subprocess.check_output(cmd, shell=True, timeout=5)
-        print('%sLogged in successfully in "%s".%s' % (GREEN, ip, DEF))
+        u.success('Logged in successfully in "%s".' % ip)
     except subprocess.CalledProcessError:
-        print('%sFailed to login using SSH, run "ssh-copy-id %s".%s' % (RED, ip, DEF))
+        u.error('Failed to login using SSH, run "ssh-copy-id %s".' % ip)
         return False
     except subprocess.TimeoutExpired:
         try:
@@ -53,26 +35,10 @@ def check_ssh(ip):
     return True
 
 
-def run_tests(ip):
-    print('Updating envsetup tests on MediaWorker.')
-    cmd = 'ssh -t root@%s /root/envsetup/update_envsetup.py' % ip
-    os.system(cmd)
-    print('Running envsetup tests on MediaWorker.')
-    cmd = 'ssh -t root@%s /root/envsetup/tester.py' % ip
-    status = os.system(cmd)
-    if status == 0:
-        subprocess.check_output(cmd, shell=True, timeout=60)
-        print('%sAll tests completed on MediaWorker "%s".' % (GREEN, ip, DEF))
-        return True
-    else:
-        print('%sapt-get update failed on MediaWorker "%s".%s' % (RED, ip, DEF))
-        return False
-
-
 def check_celerity_connectivity(ip):
-    print('Getting celerity server url.')
+    u.log('Getting celerity server url.')
     cmd = 'ssh -t root@%s cat /etc/celerity/config.py' % ip
-    print(cmd)
+    u.log(cmd)
     try:
         d = subprocess.check_output(cmd, shell=True, timeout=5, universal_newlines=True)
     except subprocess.CalledProcessError:
@@ -80,78 +46,91 @@ def check_celerity_connectivity(ip):
     else:
         m = re.search(r'\nSERVER_URL\s*=\s*([:\/\'\"\-\_\.\w]+)', d)
     if not m:
-        print('%sFailed to get celerity tasks server url from configuration in MediaWorker "%s".%s' % (RED, ip, DEF))
+        u.error('Failed to get celerity tasks server url from configuration in MediaWorker "%s".' % ip)
         return False
     server_url = m.groups()[0].strip('"\' ')
-    print('Checking celerity connectivity.')
+    u.log('Checking celerity connectivity.')
     cmd = 'ssh -t root@%s curl -k %s' % (ip, server_url)
-    print(cmd)
+    u.log(cmd)
     try:
         d = subprocess.check_output(cmd, shell=True, timeout=5, universal_newlines=True)
     except subprocess.CalledProcessError:
         d = ''
     if 'Celerity tasks server' in d:
-        print('%sSuccessfully reached tasks server from MediaWorker "%s".%s' % (GREEN, ip, DEF))
+        u.success('%sSuccessfully reached tasks server from MediaWorker "%s".%s' % ip)
         return True
-    print('%sFailed to reach tasks server from MediaWorker "%s".%s' % (RED, ip, DEF))
+    u.error('Failed to reach tasks server from MediaWorker "%s".' % ip)
     return False
 
 
 def check_celerity_versions(ip):
-    print('Checking that celerity server and worker uses the same version.')
+    u.log('Checking that celerity server and worker uses the same version.')
     try:
         ms_out = subprocess.check_output('dpkg -s celerity-utils | grep "^Version:"', shell=True, timeout=10, universal_newlines=True)
         mw_out = subprocess.check_output('ssh -t root@%s dpkg -s celerity-utils | grep "^Version:"' % ip, shell=True, timeout=10, universal_newlines=True)
     except subprocess.CalledProcessError as e:
-        print('%sFailed to check celerity version in MediaWorker "%s":\n%s%s' % (RED, ip, e, DEF))
+        u.error('Failed to check celerity version in MediaWorker "%s":\n%s' % ip, e)
         return False
     ms_out = (ms_out[len('Version:'):] if ms_out.startswith('Version:') else ms_out).strip()
     mw_out = (mw_out[len('Version:'):] if mw_out.startswith('Version:') else mw_out).strip()
     if ms_out != mw_out:
-        print('%sThe celerity version in MediaWorker "%s" is not the same as in MediaServer.%s\nMediaServer version: \tn%s\nMediaWorker version: \t%s' % (RED, ip, DEF, ms_out, mw_out))
+        u.error('The celerity version in MediaWorker "%s" is not the same as in MediaServer.\nMediaServer version: \tn%s\nMediaWorker version: \t%s' % (ip, ms_out, mw_out))
         return False
-    print('%sThe celerity version in MediaWorker "%s" is the same as in MediaServer.%s\nCurrent celerity version is: %s.' % (GREEN, ip, DEF, ms_out))
+    u.success('The celerity version in MediaWorker "%s" is the same as in MediaServer.\nCurrent celerity version is: %s.' % (ip, ms_out))
     return True
 
 
-def check_mediaworker_in_whitelist(ip):
-    # this check is not usefull anymore because the worker uses the
-    # API to get links to resources with valid secure link token
-    nginx_vhosts_path = '/etc/nginx/sites-enabled'
-    vhosts = os.listdir(nginx_vhosts_path)
-    for v in vhosts:
-        if v.endswith('.conf'):
-            vhost_path = os.path.join(nginx_vhosts_path, v)
-            with open(vhost_path, 'r') as f:
-                d = f.read()
-                if 'msuser_whitelist' in d:
-                    if ip not in d:
-                        print('%sMediaWorker ip %s is not in %s whitelist%s' % (RED, ip, v, DEF))
-                        return False
-    return True
+def run_tests(ip):
+    u.log('Updating envsetup tests on MediaWorker.')
+    cmd = 'ssh -t root@%s /root/envsetup/update_envsetup.py' % ip
+    subprocess.run(cmd, shell=True)
+    u.log('Running envsetup tests on MediaWorker.')
+    cmd = 'ssh -t root@%s /root/envsetup/tester.py' % ip
+    p = subprocess.run(cmd, shell=True)
+    if p.returncode == 0:
+        subprocess.check_output(cmd, shell=True, timeout=60)
+        u.success('All tests completed on MediaWorker "%s".' % ip)
+        return True
+    else:
+        u.error('apt-get update failed on MediaWorker "%s".' % ip)
+        return False
 
 
-all_ok = True
-tested = False
-mediaserver_ip = conf.get('NETWORK_IP')
-worker_ips = conf.get('CELERITY_WORKER_IP')
-for worker_ip in worker_ips.split(','):
-    worker_ip = worker_ip.strip()
-    if worker_ip and not worker_ip.startswith('127.0.') and worker_ip != mediaserver_ip:
-        tested = True
-        if not check_ssh(worker_ip):
-            all_ok = False
-        else:
-            if not check_celerity_connectivity(worker_ip):
-                all_ok = False
-            if not check_celerity_versions(worker_ip):
+def main():
+    try:
+        import mediaserver
+    except ImportError:
+        u.log('MediaServer is not installed, skipping test.')
+        return 2
+    else:
+        u.log('MediaServer version: %s.' % mediaserver.__version__)
+
+    all_ok = True
+    tested = False
+    mediaserver_ip = u.get_conf('NETWORK_IP')
+    worker_ips = u.get_conf('CELERITY_WORKER_IP')
+    for worker_ip in worker_ips.split(','):
+        worker_ip = worker_ip.strip()
+        if worker_ip and not worker_ip.startswith('127.0.') and worker_ip != mediaserver_ip:
+            tested = True
+            if not check_ssh(worker_ip):
                 all_ok = False
-            # if not run_tests(worker_ip):
-            #     all_ok = False
-            # if not check_mediaworker_in_whitelist(worker_ip):
-            #     all_ok = False
-if not tested:
-    print('Celerity IP not set or running locally, skipping test.')
-    sys.exit(2)
-
-sys.exit(int(not all_ok))
+            else:
+                if not check_celerity_connectivity(worker_ip):
+                    all_ok = False
+                if not check_celerity_versions(worker_ip):
+                    all_ok = False
+                # if not run_tests(worker_ip):
+                #     all_ok = False
+    if not tested:
+        u.log('Celerity IP not set or running locally, skipping test.')
+        return 2
+
+    if not all_ok:
+        return 1
+    return 0
+
+
+if __name__ == '__main__':
+    code = main()
+    sys.exit(code)

tests/test_monitoring.py

@@ -5,17 +5,15 @@ Criticality: Low
 Check that the monitoring graphs work.
 '''
 from datetime import datetime
-import imp
+from pathlib import Path
 import os
 import subprocess
 import sys
 
-os.chdir(os.path.dirname(__file__))
-if not os.path.isfile('../utils.py'):
-    print('The envsetup configuration was not found.')
-    sys.exit(1)
-u = imp.load_source('es_utils', '../utils.py')
-conf = u.load_conf()
+sys.path.append(str(Path(__file__).parents[1].resolve()))
+
+# pylint: disable=wrong-import-position
+from envsetup import utils as u  # noqa: E402
 
 
 MUNIN_WWW_PATH = '/var/cache/munin/www/'
@@ -61,5 +59,5 @@ def check_munin():
 
 
 if __name__ == '__main__':
-    rc = check_munin()
-    sys.exit(rc)
+    code = check_munin()
+    sys.exit(code)