| #!/usr/bin/python |
| |
| """Tests for drone_utility.""" |
| |
| import os, sys, unittest |
| from cStringIO import StringIO |
| |
| import common |
| from autotest_lib.client.common_lib import global_config |
| from autotest_lib.client.common_lib.test_utils import mock |
| from autotest_lib.scheduler import drone_utility |
| |
| |
| class TestDroneUtility(unittest.TestCase): |
| def setUp(self): |
| self.drone_utility = drone_utility.DroneUtility() |
| 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(self.drone_utility, '_get_process_info') |
| |
| |
| def tearDown(self): |
| self.god.unstub_all() |
| global_config.global_config.reset_config_values() |
| |
| |
| @staticmethod |
| def _set_check_dark_mark(value): |
| global_config.global_config.override_config_value( |
| 'SCHEDULER', 'check_processes_for_dark_mark', repr(value)) |
| |
| |
| def test_refresh_processes_ignore_dark_mark(self): |
| self._set_check_dark_mark(False) |
| self.drone_utility._get_process_info.expect_call().and_return( |
| [self._fake_proc_info]) |
| fake_open = lambda path, mode: self.fail('dark mark checked!') |
| processes = self.drone_utility._refresh_processes(self._fake_command, |
| open=fake_open) |
| our_pid = self._fake_proc_info['pid'] |
| for process in processes: |
| if our_pid == process['pid']: |
| break |
| else: |
| self.fail("No %s processes found" % self._fake_command) |
| self.god.check_playback() |
| |
| |
| def test_refresh_processes_check_dark_mark(self): |
| self._set_check_dark_mark(True) |
| num_procs = 2 |
| proc_info_list = num_procs * [self._fake_proc_info] |
| |
| self.drone_utility._get_process_info.expect_call().and_return( |
| proc_info_list) |
| # Test processes that have the mark in their env. |
| def _open_mark(path, mode): |
| return StringIO('foo=\0%s=\0bar=\0' % |
| drone_utility.DARK_MARK_ENVIRONMENT_VAR) |
| processes = self.drone_utility._refresh_processes(self._fake_command, |
| open=_open_mark) |
| self.assertEqual(num_procs, len(processes)) |
| self.assertEqual(proc_info_list, processes) |
| |
| self.drone_utility._get_process_info.expect_call().and_return( |
| proc_info_list) |
| # Test processes that do not have the mark in their env |
| def _open_nomark(path, mode): |
| return StringIO('foo=\0bar=\0') # No dark mark. |
| processes = self.drone_utility._refresh_processes(self._fake_command, |
| open=_open_nomark) |
| self.assertEqual([], processes) |
| self.god.check_playback() |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
