scheduler/drone_utility_unittest.py

#!/usr/bin/python2

"""Tests for drone_utility."""

import os
import unittest

import common
from autotest_lib.client.common_lib import autotemp
from autotest_lib.client.common_lib.test_utils import mock
from autotest_lib.scheduler import drone_utility


class TestProcessRefresher(unittest.TestCase):
    """Tests for the drone_utility.ProcessRefresher object."""

    def setUp(self):
        self._tempdir = autotemp.tempdir(unique_id='test_process_refresher')
        self.addCleanup(self._tempdir.clean)
        self._fake_command = '!faketest!'
        self._fake_proc_info = {'pid': 3, 'pgid': 4, 'ppid': 2,
                                'comm': self._fake_command, 'args': ''}
        self.god = mock.mock_god()
        self.god.stub_function(drone_utility, '_get_process_info')
        self._mock_get_process_info = drone_utility._get_process_info
        self.god.stub_function(drone_utility, '_process_has_dark_mark')
        self._mock_process_has_dark_mark = (
                drone_utility._process_has_dark_mark)


    def tearDown(self):
        self.god.unstub_all()

    def test_no_processes(self):
        """Sanity check the case when there is nothing to do"""
        self._mock_get_process_info.expect_call().and_return([])
        process_refresher = drone_utility.ProcessRefresher(check_mark=False)
        got, warnings = process_refresher([])
        expected = {
                'pidfiles': dict(),
                'all_processes': [],
                'autoserv_processes': [],
                'parse_processes': [],
                'pidfiles_second_read': dict(),
        }
        self.god.check_playback()
        self.assertEqual(got, expected)


    def test_read_pidfiles_use_pool(self):
        """Readable subset of pidfile paths are included in the result

        Uses process pools.
        """
        self._parameterized_test_read_pidfiles(use_pool=True)


    def test_read_pidfiles_no_pool(self):
        """Readable subset of pidfile paths are included in the result

        Does not use process pools.
        """
        self._parameterized_test_read_pidfiles(use_pool=False)


    def test_read_many_pidfiles(self):
        """Read a large number of pidfiles (more than pool size)."""
        self._mock_get_process_info.expect_call().and_return([])
        expected_pidfiles = {}
        for i in range(1000):
            data = 'data number %d' % i
            path = self._write_pidfile('pidfile%d' % i, data)
            expected_pidfiles[path] = data

        process_refresher = drone_utility.ProcessRefresher(check_mark=False,
                                                           use_pool=True)
        got, _ = process_refresher(expected_pidfiles.keys())
        expected = {
                'pidfiles': expected_pidfiles,
                'all_processes': [],
                'autoserv_processes': [],
                'parse_processes': [],
                'pidfiles_second_read': expected_pidfiles,
        }
        self.god.check_playback()
        self.assertEqual(got, expected)


    def test_filter_processes(self):
        """Various filtered results correctly classify processes by name."""
        self.maxDiff = None
        process_refresher = drone_utility.ProcessRefresher(check_mark=False)
        autoserv_processes = [self._proc_info_dict(3, 'autoserv')]
        parse_processes = [self._proc_info_dict(4, 'parse'),
                           self._proc_info_dict(5, 'site_parse')]
        all_processes = ([self._proc_info_dict(6, 'who_cares')]
                         + autoserv_processes + parse_processes)

        self._mock_get_process_info.expect_call().and_return(all_processes)
        got, _warnings = process_refresher(self._tempdir.name)
        expected = {
                'pidfiles': dict(),
                'all_processes': all_processes,
                'autoserv_processes': autoserv_processes,
                'parse_processes': parse_processes,
                'pidfiles_second_read': dict(),
        }
        self.god.check_playback()
        self.assertEqual(got, expected)


    def test_respect_dark_mark(self):
        """When check_mark=True, dark mark check is performed and respected.

        Only filtered processes with dark mark should be returned. We only test
        this with use_pool=False because mocking out _process_has_dark_mark with
        multiprocessing.Pool is hard.
        """
        self.maxDiff = None
        process_refresher = drone_utility.ProcessRefresher(check_mark=True)
        marked_process = self._proc_info_dict(3, 'autoserv')
        unmarked_process = self._proc_info_dict(369, 'autoserv')
        all_processes = [marked_process, unmarked_process]
        self._mock_get_process_info.expect_call().and_return(all_processes)
        self._mock_process_has_dark_mark.expect_call(3).and_return(True)
        self._mock_process_has_dark_mark.expect_call(369).and_return(False)
        got, warnings = process_refresher(self._tempdir.name)
        expected = {
                'pidfiles': dict(),
                'all_processes': all_processes,
                'autoserv_processes': [marked_process],
                'parse_processes': [],
                'pidfiles_second_read': dict(),
        }
        self.god.check_playback()
        self.assertEqual(got, expected)
        self.assertEqual(len(warnings), 1)
        self.assertRegexpMatches(warnings[0], '.*autoserv.*369.*')


    def _parameterized_test_read_pidfiles(self, use_pool):
        """Readable subset of pidfile paths are included in the result

        @param: use_pool: Argument use_pool for ProcessRefresher
        """
        self._mock_get_process_info.expect_call().and_return([])
        path1 = self._write_pidfile('pidfile1', 'first pidfile')
        path2 = self._write_pidfile('pidfile2', 'second pidfile',
                                    subdir='somedir')
        process_refresher = drone_utility.ProcessRefresher(check_mark=False,
                                                           use_pool=use_pool)
        got, warnings = process_refresher(
                [path1, path2,
                 os.path.join(self._tempdir.name, 'non_existent')])
        expected_pidfiles = {
                path1: 'first pidfile',
                path2: 'second pidfile',
        }
        expected = {
                'pidfiles': expected_pidfiles,
                'all_processes': [],
                'autoserv_processes': [],
                'parse_processes': [],
                'pidfiles_second_read': expected_pidfiles,
        }
        self.god.check_playback()
        self.assertEqual(got, expected)


    def _write_pidfile(self, filename, content, subdir=''):
        parent_dir = self._tempdir.name
        if subdir:
            parent_dir = os.path.join(parent_dir, subdir)
            os.makedirs(parent_dir)
        path = os.path.join(parent_dir, filename)
        with open(path, 'w') as f:
            f.write(content)
        return path

    def _proc_info_dict(self, pid, comm, pgid=33, ppid=44, args=''):
        return {'pid': pid, 'comm': comm, 'pgid': pgid, 'ppid': ppid,
                'args': args}


if __name__ == '__main__':
    unittest.main()