Use the already defined function to wait for generate

[pyi2ncommon] / test / test_arnied_wrapper.py

#!/usr/bin/env python

# The software in this package is distributed under the GNU General
# Public License version 2 (with a special exception described below).
#
# A copy of GNU General Public License (GPL) is included in this distribution,
# in the file COPYING.GPL.
#
# As a special exception, if other files instantiate templates or use macros
# or inline functions from this file, or you compile this file and link it
# with other works to produce a work based on this file, this file
# does not by itself cause the resulting work to be covered
# by the GNU General Public License.
#
# However the source code for this file must still be made available
# in accordance with section (3) of the GNU General Public License.
#
# This exception does not invalidate any other reasons why a work based
# on this file might be covered by the GNU General Public License.
#
# Copyright (c) 2016-2018 Intra2net AG <info@intra2net.com>

import unittest
import unittest.mock as mock
import subprocess

from src import arnied_wrapper


class DummyCmdOutputMapping(object):
    """
    Class to replace the :py:function:`arnied_wrapper.run_cmd` function.

    In the arnied_wrapper, when running a command, instead of calling the actual
    function, this class' constructor will be invoked and an instance returned
    instead. It stubs :py:class:`subprocess.CompletedProcess`, such that the
    returncode and stdout attributes are set depending on the command or on the
    test.
    """

    # whether to return 1 as a fail indicator
    fail_switch = False
    # mapping between expected commands and their mocked output + return code
    cmds = [
        {"cmd": "pgrep -l -x arnied", "stdout": b"", "returncode": 0},

        {"cmd": 'echo "LICENSE_ACCEPTED,0: \\"1\\"" | set_cnf', "stdout": b"", "returncode": 0},
        {"cmd": '/usr/intranator/bin/arnied_helper --wait-for-program-end GENERATE', "stdout": b"", "returncode": 0},

        {"cmd": 'get_cnf PROVIDER 1', "stdout": b"1  PROVIDER,1: \"sample-provider\"", "returncode": 0},
        {"cmd": 'tell-connd --online P1', "stdout": b"", "returncode": 0},
        {"cmd": '/usr/intranator/bin/get_var ONLINE', "stdout": b"DEFAULT: 2", "returncode": 0},

        {"cmd": "pgrep -l -x arnied", "stdout": b"", "returncode": 0},
        {"cmd": 'get_cnf VIRSCAN_UPDATE_CRON  | set_cnf -x', "stdout": b"", "returncode": 0},
        {"cmd": '/usr/intranator/bin/arnied_helper --is-scheduled-or-running GENERATE', "stdout": b"", "returncode": 1},
        {"cmd": '/usr/intranator/bin/arnied_helper --wait-for-program-end GENERATE --wait-for-program-timeout 300', "stdout": b"", "returncode": 1},
        {"cmd": '/usr/intranator/bin/arnied_helper --is-scheduled-or-running GENERATE_OFFLINE', "stdout": b"", "returncode": 1},
        {"cmd": '/usr/intranator/bin/arnied_helper --wait-for-program-end GENERATE_OFFLINE --wait-for-program-timeout 300', "stdout": b"", "returncode": 1},
        {"cmd": 'echo \'VIRSCAN_UPDATE_DNS_PUSH,0:"0"\' |set_cnf', "stdout": b"", "returncode": 0},
        {"cmd": 'rm -f /var/intranator/schedule/UPDATE_VIRSCAN_NODIAL*', "stdout": b"", "returncode": 0},

        {"cmd": '/usr/intranator/bin/arnied_helper --transfer-mail', "stdout": b"", "returncode": 0},

        {"cmd": '/usr/intranator/bin/arnied_helper --wait-for-arnied-socket --wait-for-arnied-socket-timeout 10', "stdout": b"", "returncode": 0},
        {"cmd": '/usr/intranator/bin/arnied_helper --wait-for-arnied-socket --wait-for-arnied-socket-timeout 30', "stdout": b"", "returncode": 0},
    ]
    asserted_cmds = []

    def __init__(self, cmd="", ignore_errors=False, vm=None, timeout=60):
        """
        Class constructor to mimic the run function of the arnied wrapper.

        Arguments are the same of the mocked function.
        """
        self.returncode = 0
        self.stdout = b""
        self._get_result(cmd)
        if self.fail_switch:
            self.returncode = 1

    def __str__(self):
        """String representation of this class."""
        return "status %i, stdout %s" % (self.returncode, self.stdout)

    def _get_result(self, cmd):
        """
        Return the first (or raise) values in the mapping matching the command.

        :param str cmd: command to check the mapping against
        :raises: :py:class:`ValueError` if the command has no corresponding
                 mapping
        :returns: this instance with the return code and stdout attributes set
        :rtype: :py:class:`DummyCmdOutputMapping`
        """
        for dummy_cmd in self.asserted_cmds:
            if dummy_cmd['cmd'] == cmd:
                self.returncode = dummy_cmd['returncode']
                self.stdout = dummy_cmd['stdout']
                return
        raise ValueError("Could not locate the command '%s' among the known answers "
                         "for the universe" % cmd)


# make sure that invoking `run_cmd` returns an instance of DummyCmdOutputMapping
@mock.patch("src.arnied_wrapper.run_cmd", DummyCmdOutputMapping)
class ArniedWrapperTest(unittest.TestCase):
    def setUp(self):
        DummyCmdOutputMapping.fail_switch = False
        DummyCmdOutputMapping.asserted_cmds = []
        self.cmd_db = DummyCmdOutputMapping.cmds

    def test_verify_running(self):
        """Test checking for running programs."""
        DummyCmdOutputMapping.asserted_cmds = self.cmd_db[0:1]
        arnied_wrapper.verify_running(timeout=1)
        DummyCmdOutputMapping.fail_switch = True
        with self.assertRaises(RuntimeError):
            arnied_wrapper.verify_running(timeout=1)

    def test_wait_for_arnied(self):
        """Test waiting for arnied to be ready."""
        DummyCmdOutputMapping.asserted_cmds = self.cmd_db
        arnied_wrapper.wait_for_arnied(timeout=10)

    def test_accept_license(self):
        """Test accepting license."""
        DummyCmdOutputMapping.asserted_cmds = self.cmd_db[1:3] + self.cmd_db[8:11]
        arnied_wrapper.accept_licence()
        # make sure an error is ignored since license might
        # already be accepted
        DummyCmdOutputMapping.fail_switch = True
        arnied_wrapper.accept_licence()

    def test_go_online(self):
        DummyCmdOutputMapping.asserted_cmds = self.cmd_db[3:6]
        arnied_wrapper.go_online(1)

    def test_disable_virscan(self):
        """Test disabling the virus scanner."""
        DummyCmdOutputMapping.asserted_cmds = self.cmd_db[6:]
        arnied_wrapper.disable_virscan()

    def test_email_transfer(self):
        """Test e-mail transferring."""
        DummyCmdOutputMapping.asserted_cmds = self.cmd_db[14:15]
        arnied_wrapper.email_transfer()