scheduler/drone_utility_unittest.py

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