1#!/usr/bin/python
2
3"""Tests for drone_utility."""
4
5import os, sys, unittest
6from cStringIO import StringIO
7
8import common
9from autotest_lib.client.common_lib import global_config
10from autotest_lib.client.common_lib.test_utils import mock
11from autotest_lib.scheduler import drone_utility
12
13
14class 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
78if __name__ == '__main__':
79    unittest.main()
80