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() 80
