Home | History | Annotate | Download | only in scheduler 
      1 #!/usr/bin/python
      2 
      3 """Tests for drone_utility."""
      4 
      5 import os
      6 import unittest
      7 
      8 import common
      9 from autotest_lib.client.common_lib import autotemp
     10 from autotest_lib.client.common_lib.test_utils import mock
     11 from autotest_lib.scheduler import drone_utility
     12 
     13 
     14 class TestProcessRefresher(unittest.TestCase):
     15     """Tests for the drone_utility.ProcessRefresher object."""
     16 
     17     def setUp(self):
     18         self._tempdir = autotemp.tempdir(unique_id='test_process_refresher')
     19         self.addCleanup(self._tempdir.clean)
     20         self._fake_command = '!faketest!'
     21         self._fake_proc_info = {'pid': 3, 'pgid': 4, 'ppid': 2,
     22                                 'comm': self._fake_command, 'args': ''}
     23         self.god = mock.mock_god()
     24         self.god.stub_function(drone_utility, '_get_process_info')
     25         self._mock_get_process_info = drone_utility._get_process_info
     26         self.god.stub_function(drone_utility, '_process_has_dark_mark')
     27         self._mock_process_has_dark_mark = (
     28                 drone_utility._process_has_dark_mark)
     29 
     30 
     31     def tearDown(self):
     32         self.god.unstub_all()
     33 
     34     def test_no_processes(self):
     35         """Sanity check the case when there is nothing to do"""
     36         self._mock_get_process_info.expect_call().and_return([])
     37         process_refresher = drone_utility.ProcessRefresher(check_mark=False)
     38         got, warnings = process_refresher([])
     39         expected = {
     40                 'pidfiles': dict(),
     41                 'all_processes': [],
     42                 'autoserv_processes': [],
     43                 'parse_processes': [],
     44                 'pidfiles_second_read': dict(),
     45         }
     46         self.god.check_playback()
     47         self.assertEqual(got, expected)
     48 
     49 
     50     def test_read_pidfiles_use_pool(self):
     51         """Readable subset of pidfile paths are included in the result
     52 
     53         Uses process pools.
     54         """
     55         self._parameterized_test_read_pidfiles(use_pool=True)
     56 
     57 
     58     def test_read_pidfiles_no_pool(self):
     59         """Readable subset of pidfile paths are included in the result
     60 
     61         Does not use process pools.
     62         """
     63         self._parameterized_test_read_pidfiles(use_pool=False)
     64 
     65 
     66     def test_read_many_pidfiles(self):
     67         """Read a large number of pidfiles (more than pool size)."""
     68         self._mock_get_process_info.expect_call().and_return([])
     69         expected_pidfiles = {}
     70         for i in range(1000):
     71             data = 'data number %d' % i
     72             path = self._write_pidfile('pidfile%d' % i, data)
     73             expected_pidfiles[path] = data
     74 
     75         process_refresher = drone_utility.ProcessRefresher(check_mark=False,
     76                                                            use_pool=True)
     77         got, _ = process_refresher(expected_pidfiles.keys())
     78         expected = {
     79                 'pidfiles': expected_pidfiles,
     80                 'all_processes': [],
     81                 'autoserv_processes': [],
     82                 'parse_processes': [],
     83                 'pidfiles_second_read': expected_pidfiles,
     84         }
     85         self.god.check_playback()
     86         self.assertEqual(got, expected)
     87 
     88 
     89     def test_filter_processes(self):
     90         """Various filtered results correctly classify processes by name."""
     91         self.maxDiff = None
     92         process_refresher = drone_utility.ProcessRefresher(check_mark=False)
     93         autoserv_processes = [self._proc_info_dict(3, 'autoserv')]
     94         parse_processes = [self._proc_info_dict(4, 'parse'),
     95                            self._proc_info_dict(5, 'site_parse')]
     96         all_processes = ([self._proc_info_dict(6, 'who_cares')]
     97                          + autoserv_processes + parse_processes)
     98 
     99         self._mock_get_process_info.expect_call().and_return(all_processes)
    100         got, _warnings = process_refresher(self._tempdir.name)
    101         expected = {
    102                 'pidfiles': dict(),
    103                 'all_processes': all_processes,
    104                 'autoserv_processes': autoserv_processes,
    105                 'parse_processes': parse_processes,
    106                 'pidfiles_second_read': dict(),
    107         }
    108         self.god.check_playback()
    109         self.assertEqual(got, expected)
    110 
    111 
    112     def test_respect_dark_mark(self):
    113         """When check_mark=True, dark mark check is performed and respected.
    114 
    115         Only filtered processes with dark mark should be returned. We only test
    116         this with use_pool=False because mocking out _process_has_dark_mark with
    117         multiprocessing.Pool is hard.
    118         """
    119         self.maxDiff = None
    120         process_refresher = drone_utility.ProcessRefresher(check_mark=True)
    121         marked_process = self._proc_info_dict(3, 'autoserv')
    122         unmarked_process = self._proc_info_dict(369, 'autoserv')
    123         all_processes = [marked_process, unmarked_process]
    124         self._mock_get_process_info.expect_call().and_return(all_processes)
    125         self._mock_process_has_dark_mark.expect_call(3).and_return(True)
    126         self._mock_process_has_dark_mark.expect_call(369).and_return(False)
    127         got, warnings = process_refresher(self._tempdir.name)
    128         expected = {
    129                 'pidfiles': dict(),
    130                 'all_processes': all_processes,
    131                 'autoserv_processes': [marked_process],
    132                 'parse_processes': [],
    133                 'pidfiles_second_read': dict(),
    134         }
    135         self.god.check_playback()
    136         self.assertEqual(got, expected)
    137         self.assertEqual(len(warnings), 1)
    138         self.assertRegexpMatches(warnings[0], '.*autoserv.*369.*')
    139 
    140 
    141     def _parameterized_test_read_pidfiles(self, use_pool):
    142         """Readable subset of pidfile paths are included in the result
    143 
    144         @param: use_pool: Argument use_pool for ProcessRefresher
    145         """
    146         self._mock_get_process_info.expect_call().and_return([])
    147         path1 = self._write_pidfile('pidfile1', 'first pidfile')
    148         path2 = self._write_pidfile('pidfile2', 'second pidfile',
    149                                     subdir='somedir')
    150         process_refresher = drone_utility.ProcessRefresher(check_mark=False,
    151                                                            use_pool=use_pool)
    152         got, warnings = process_refresher(
    153                 [path1, path2,
    154                  os.path.join(self._tempdir.name, 'non_existent')])
    155         expected_pidfiles = {
    156                 path1: 'first pidfile',
    157                 path2: 'second pidfile',
    158         }
    159         expected = {
    160                 'pidfiles': expected_pidfiles,
    161                 'all_processes': [],
    162                 'autoserv_processes': [],
    163                 'parse_processes': [],
    164                 'pidfiles_second_read': expected_pidfiles,
    165         }
    166         self.god.check_playback()
    167         self.assertEqual(got, expected)
    168 
    169 
    170     def _write_pidfile(self, filename, content, subdir=''):
    171         parent_dir = self._tempdir.name
    172         if subdir:
    173             parent_dir = os.path.join(parent_dir, subdir)
    174             os.makedirs(parent_dir)
    175         path = os.path.join(parent_dir, filename)
    176         with open(path, 'w') as f:
    177             f.write(content)
    178         return path
    179 
    180     def _proc_info_dict(self, pid, comm, pgid=33, ppid=44, args=''):
    181         return {'pid': pid, 'comm': comm, 'pgid': pgid, 'ppid': ppid,
    182                 'args': args}
    183 
    184 
    185 if __name__ == '__main__':
    186     unittest.main()
    187