test_mediaworker.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Criticality: High
Checks that MediaWorker can be reached using SSH and that it can reach the tasks server
'''
from pathlib import Path
import re
import subprocess
import sys

sys.path.append(str(Path(__file__).parents[1].resolve()))

# 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
    u.log('Connecting to MediaWorker:\n%s' % cmd)
    try:
        subprocess.check_output(cmd, shell=True, timeout=5)
        u.success('Logged in successfully in "%s".' % ip)
    except subprocess.CalledProcessError:
        u.error('Failed to login using SSH, run "ssh-copy-id %s".' % ip)
        return False
    except subprocess.TimeoutExpired:
        try:
            cmd_port = 'nc -z -w2 {} 22'.format(ip)
            subprocess.check_output(cmd_port, shell=True, timeout=5)
        except Exception:
            cmd_ping = 'ping -c2 -w4 {}'.format(ip)
            subprocess.check_output(cmd_ping, shell=True, timeout=15)
        return False
    return True


def check_celerity_connectivity(ip):
    u.log('Getting celerity server url.')
    cmd = 'ssh -t root@%s cat /etc/celerity/config.py' % ip
    u.log(cmd)
    try:
        d = subprocess.check_output(cmd, shell=True, timeout=5, universal_newlines=True)
    except subprocess.CalledProcessError:
        m = None
    else:
        m = re.search(r'\nSERVER_URL\s*=\s*([:\/\'\"\-\_\.\w]+)', d)
    if not m:
        u.error('Failed to get celerity tasks server url from configuration in MediaWorker "%s".' % ip)
        return False
    server_url = m.groups()[0].strip('"\' ')
    u.log('Checking celerity connectivity.')
    cmd = 'ssh -t root@%s curl -k %s' % (ip, server_url)
    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:
        u.success('Successfully reached tasks server from MediaWorker "%s".' % ip)
        return True
    u.error('Failed to reach tasks server from MediaWorker "%s".' % ip)
    return False


def check_celerity_versions(ip):
    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:
        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:
        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
    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 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


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
            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)