showard | d791dcb | 2009-09-16 17:17:36 +0000 | [diff] [blame] | 1 | #!/usr/bin/python |
| 2 | |
| 3 | """Tests for drone_utility.""" |
| 4 | |
| 5 | import os, sys, unittest |
| 6 | from cStringIO import StringIO |
| 7 | |
| 8 | import common |
| 9 | from autotest_lib.client.common_lib import global_config |
| 10 | from autotest_lib.client.common_lib.test_utils import mock |
| 11 | from autotest_lib.scheduler import drone_utility |
| 12 | |
| 13 | |
| 14 | class TestDroneUtility(unittest.TestCase): |
| 15 | def setUp(self): |
| 16 | self.drone_utility = drone_utility.DroneUtility() |
| 17 | self._fake_command = '!faketest!' |
| 18 | self._fake_proc_info = {'pid': 3, 'pgid': 4, 'ppid': 2, |
| 19 | 'comm': self._fake_command, 'args': ''} |
| 20 | self.god = mock.mock_god() |
| 21 | self.god.stub_function(self.drone_utility, '_get_process_info') |
| 22 | |
| 23 | |
| 24 | def tearDown(self): |
| 25 | self.god.unstub_all() |
| 26 | global_config.global_config.reset_config_values() |
| 27 | |
| 28 | |
| 29 | @staticmethod |
| 30 | def _set_check_dark_mark(value): |
| 31 | global_config.global_config.override_config_value( |
| 32 | 'SCHEDULER', 'check_processes_for_dark_mark', repr(value)) |
| 33 | |
| 34 | |
| 35 | def test_refresh_processes_ignore_dark_mark(self): |
| 36 | self._set_check_dark_mark(False) |
| 37 | self.drone_utility._get_process_info.expect_call().and_return( |
| 38 | [self._fake_proc_info]) |
| 39 | fake_open = lambda path, mode: self.fail('dark mark checked!') |
| 40 | processes = self.drone_utility._refresh_processes(self._fake_command, |
| 41 | open=fake_open) |
| 42 | our_pid = self._fake_proc_info['pid'] |
| 43 | for process in processes: |
| 44 | if our_pid == process['pid']: |
| 45 | break |
| 46 | else: |
| 47 | self.fail("No %s processes found" % self._fake_command) |
| 48 | self.god.check_playback() |
| 49 | |
| 50 | |
| 51 | def test_refresh_processes_check_dark_mark(self): |
| 52 | self._set_check_dark_mark(True) |
| 53 | num_procs = 2 |
| 54 | proc_info_list = num_procs * [self._fake_proc_info] |
| 55 | |
| 56 | self.drone_utility._get_process_info.expect_call().and_return( |
| 57 | proc_info_list) |
| 58 | # Test processes that have the mark in their env. |
| 59 | def _open_mark(path, mode): |
| 60 | return StringIO('foo=\0%s=\0bar=\0' % |
| 61 | drone_utility.DARK_MARK_ENVIRONMENT_VAR) |
| 62 | processes = self.drone_utility._refresh_processes(self._fake_command, |
| 63 | open=_open_mark) |
| 64 | self.assertEqual(num_procs, len(processes)) |
| 65 | self.assertEqual(proc_info_list, processes) |
| 66 | |
| 67 | self.drone_utility._get_process_info.expect_call().and_return( |
| 68 | proc_info_list) |
| 69 | # Test processes that do not have the mark in their env |
| 70 | def _open_nomark(path, mode): |
| 71 | return StringIO('foo=\0bar=\0') # No dark mark. |
| 72 | processes = self.drone_utility._refresh_processes(self._fake_command, |
| 73 | open=_open_nomark) |
| 74 | self.assertEqual([], processes) |
| 75 | self.god.check_playback() |
| 76 | |
| 77 | |
| 78 | if __name__ == '__main__': |
| 79 | unittest.main() |