Home | History | Annotate | Download | only in server 
      1 #!/usr/bin/python
      2 # Copyright 2009 Google Inc. Released under the GPL v2
      3 
      4 import time, unittest
      5 
      6 import common
      7 from autotest_lib.client.common_lib import error
      8 from autotest_lib.client.common_lib.test_utils import mock
      9 from autotest_lib.server import subcommand
     10 
     11 
     12 def _create_subcommand(func, args):
     13     # to avoid __init__
     14     class wrapper(subcommand.subcommand):
     15         def __init__(self, func, args):
     16             self.func = func
     17             self.args = args
     18             self.subdir = None
     19             self.debug = None
     20             self.pid = None
     21             self.returncode = None
     22 
     23     return wrapper(func, args)
     24 
     25 
     26 class subcommand_test(unittest.TestCase):
     27     def setUp(self):
     28         self.god = mock.mock_god()
     29 
     30 
     31     def tearDown(self):
     32         self.god.unstub_all()
     33         # cleanup the hooks
     34         subcommand.subcommand.fork_hooks = []
     35         subcommand.subcommand.join_hooks = []
     36 
     37 
     38     def test_create(self):
     39         def check_attributes(cmd, func, args, subdir=None, debug=None,
     40                              pid=None, returncode=None, fork_hooks=[],
     41                              join_hooks=[]):
     42             self.assertEquals(cmd.func, func)
     43             self.assertEquals(cmd.args, args)
     44             self.assertEquals(cmd.subdir, subdir)
     45             self.assertEquals(cmd.debug, debug)
     46             self.assertEquals(cmd.pid, pid)
     47             self.assertEquals(cmd.returncode, returncode)
     48             self.assertEquals(cmd.fork_hooks, fork_hooks)
     49             self.assertEquals(cmd.join_hooks, join_hooks)
     50 
     51         def func(arg1, arg2):
     52             pass
     53 
     54         cmd = subcommand.subcommand(func, (2, 3))
     55         check_attributes(cmd, func, (2, 3))
     56         self.god.check_playback()
     57 
     58         self.god.stub_function(subcommand.os.path, 'abspath')
     59         self.god.stub_function(subcommand.os.path, 'exists')
     60         self.god.stub_function(subcommand.os, 'mkdir')
     61 
     62         subcommand.os.path.abspath.expect_call('dir').and_return('/foo/dir')
     63         subcommand.os.path.exists.expect_call('/foo/dir').and_return(False)
     64         subcommand.os.mkdir.expect_call('/foo/dir')
     65 
     66         (subcommand.os.path.exists.expect_call('/foo/dir/debug')
     67                 .and_return(False))
     68         subcommand.os.mkdir.expect_call('/foo/dir/debug')
     69 
     70         cmd = subcommand.subcommand(func, (2, 3), subdir='dir')
     71         check_attributes(cmd, func, (2, 3), subdir='/foo/dir',
     72                          debug='/foo/dir/debug')
     73         self.god.check_playback()
     74 
     75 
     76     def _setup_fork_start_parent(self):
     77         self.god.stub_function(subcommand.os, 'fork')
     78 
     79         subcommand.os.fork.expect_call().and_return(1000)
     80         func = self.god.create_mock_function('func')
     81         cmd = _create_subcommand(func, [])
     82         cmd.fork_start()
     83 
     84         return cmd
     85 
     86 
     87     def test_fork_start_parent(self):
     88         cmd = self._setup_fork_start_parent()
     89 
     90         self.assertEquals(cmd.pid, 1000)
     91         self.god.check_playback()
     92 
     93 
     94     def _setup_fork_start_child(self):
     95         self.god.stub_function(subcommand.os, 'pipe')
     96         self.god.stub_function(subcommand.os, 'fork')
     97         self.god.stub_function(subcommand.os, 'close')
     98         self.god.stub_function(subcommand.os, 'write')
     99         self.god.stub_function(subcommand.cPickle, 'dumps')
    100         self.god.stub_function(subcommand.os, '_exit')
    101 
    102 
    103     def test_fork_start_child(self):
    104         self._setup_fork_start_child()
    105 
    106         func = self.god.create_mock_function('func')
    107         fork_hook = self.god.create_mock_function('fork_hook')
    108         join_hook = self.god.create_mock_function('join_hook')
    109 
    110         subcommand.subcommand.register_fork_hook(fork_hook)
    111         subcommand.subcommand.register_join_hook(join_hook)
    112         cmd = _create_subcommand(func, (1, 2))
    113 
    114         subcommand.os.pipe.expect_call().and_return((10, 20))
    115         subcommand.os.fork.expect_call().and_return(0)
    116         subcommand.os.close.expect_call(10)
    117         fork_hook.expect_call(cmd)
    118         func.expect_call(1, 2).and_return(True)
    119         subcommand.cPickle.dumps.expect_call(True,
    120                 subcommand.cPickle.HIGHEST_PROTOCOL).and_return('True')
    121         subcommand.os.write.expect_call(20, 'True')
    122         subcommand.os.close.expect_call(20)
    123         join_hook.expect_call(cmd)
    124         subcommand.os._exit.expect_call(0)
    125 
    126         cmd.fork_start()
    127         self.god.check_playback()
    128 
    129 
    130     def test_fork_start_child_error(self):
    131         self._setup_fork_start_child()
    132         self.god.stub_function(subcommand.logging, 'exception')
    133 
    134         func = self.god.create_mock_function('func')
    135         cmd = _create_subcommand(func, (1, 2))
    136         error = Exception('some error')
    137 
    138         subcommand.os.pipe.expect_call().and_return((10, 20))
    139         subcommand.os.fork.expect_call().and_return(0)
    140         subcommand.os.close.expect_call(10)
    141         func.expect_call(1, 2).and_raises(error)
    142         subcommand.logging.exception.expect_call('function failed')
    143         subcommand.cPickle.dumps.expect_call(error,
    144                 subcommand.cPickle.HIGHEST_PROTOCOL).and_return('error')
    145         subcommand.os.write.expect_call(20, 'error')
    146         subcommand.os.close.expect_call(20)
    147         subcommand.os._exit.expect_call(1)
    148 
    149         cmd.fork_start()
    150         self.god.check_playback()
    151 
    152 
    153     def _setup_poll(self):
    154         cmd = self._setup_fork_start_parent()
    155         self.god.stub_function(subcommand.os, 'waitpid')
    156         return cmd
    157 
    158 
    159     def test_poll_running(self):
    160         cmd = self._setup_poll()
    161 
    162         (subcommand.os.waitpid.expect_call(1000, subcommand.os.WNOHANG)
    163                 .and_raises(subcommand.os.error('waitpid')))
    164         self.assertEquals(cmd.poll(), None)
    165         self.god.check_playback()
    166 
    167 
    168     def test_poll_finished_success(self):
    169         cmd = self._setup_poll()
    170 
    171         (subcommand.os.waitpid.expect_call(1000, subcommand.os.WNOHANG)
    172                 .and_return((1000, 0)))
    173         self.assertEquals(cmd.poll(), 0)
    174         self.god.check_playback()
    175 
    176 
    177     def test_poll_finished_failure(self):
    178         cmd = self._setup_poll()
    179         self.god.stub_function(cmd, '_handle_exitstatus')
    180 
    181         (subcommand.os.waitpid.expect_call(1000, subcommand.os.WNOHANG)
    182                 .and_return((1000, 10)))
    183         cmd._handle_exitstatus.expect_call(10).and_raises(Exception('fail'))
    184 
    185         self.assertRaises(Exception, cmd.poll)
    186         self.god.check_playback()
    187 
    188 
    189     def test_wait_success(self):
    190         cmd = self._setup_poll()
    191 
    192         (subcommand.os.waitpid.expect_call(1000, 0)
    193                 .and_return((1000, 0)))
    194 
    195         self.assertEquals(cmd.wait(), 0)
    196         self.god.check_playback()
    197 
    198 
    199     def test_wait_failure(self):
    200         cmd = self._setup_poll()
    201         self.god.stub_function(cmd, '_handle_exitstatus')
    202 
    203         (subcommand.os.waitpid.expect_call(1000, 0)
    204                 .and_return((1000, 10)))
    205 
    206         cmd._handle_exitstatus.expect_call(10).and_raises(Exception('fail'))
    207         self.assertRaises(Exception, cmd.wait)
    208         self.god.check_playback()
    209 
    210 
    211 class real_subcommand_test(unittest.TestCase):
    212     """Test actually running subcommands (without mocking)."""
    213 
    214 
    215     def _setup_subcommand(self, func, *args):
    216         cmd = subcommand.subcommand(func, args)
    217         cmd.fork_start()
    218         return cmd
    219 
    220 
    221     def test_fork_waitfor_no_timeout(self):
    222         """Test fork_waitfor success with no timeout."""
    223         cmd = self._setup_subcommand(lambda: None)
    224         self.assertEquals(cmd.fork_waitfor(), 0)
    225 
    226 
    227     def test_fork_waitfor_timeout(self):
    228         """Test fork_waitfor success with a timeout."""
    229         cmd = self._setup_subcommand(lambda: None)
    230         self.assertEquals(cmd.fork_waitfor(timeout=60), 0)
    231 
    232 
    233     def test_fork_waitfor_exception(self):
    234         """Test fork_waitfor failure with an exception."""
    235         cmd = self._setup_subcommand(lambda: None, 'foo')
    236         with self.assertRaises(error.AutoservSubcommandError):
    237           cmd.fork_waitfor(timeout=60)
    238 
    239 
    240     def test_fork_waitfor_timeout_fail(self):
    241         """Test fork_waitfor timing out."""
    242         cmd = self._setup_subcommand(lambda: time.sleep(60))
    243         with self.assertRaises(error.AutoservSubcommandError):
    244           cmd.fork_waitfor(timeout=1)
    245 
    246 
    247 class parallel_test(unittest.TestCase):
    248     def setUp(self):
    249         self.god = mock.mock_god()
    250         self.god.stub_function(subcommand.cPickle, 'load')
    251 
    252 
    253     def tearDown(self):
    254         self.god.unstub_all()
    255 
    256 
    257     def _get_cmd(self, func, args):
    258         cmd = _create_subcommand(func, args)
    259         cmd.result_pickle = self.god.create_mock_class(file, 'file')
    260         return self.god.create_mock_class(cmd, 'subcommand')
    261 
    262 
    263     def _get_tasklist(self):
    264         return [self._get_cmd(lambda x: x * 2, (3,)),
    265                 self._get_cmd(lambda: None, [])]
    266 
    267 
    268     def _setup_common(self):
    269         tasklist = self._get_tasklist()
    270 
    271         for task in tasklist:
    272             task.fork_start.expect_call()
    273 
    274         return tasklist
    275 
    276 
    277     def test_success(self):
    278         tasklist = self._setup_common()
    279 
    280         for task in tasklist:
    281             task.fork_waitfor.expect_call(timeout=None).and_return(0)
    282             (subcommand.cPickle.load.expect_call(task.result_pickle)
    283                     .and_return(6))
    284             task.result_pickle.close.expect_call()
    285 
    286         subcommand.parallel(tasklist)
    287         self.god.check_playback()
    288 
    289 
    290     def test_failure(self):
    291         tasklist = self._setup_common()
    292 
    293         for task in tasklist:
    294             task.fork_waitfor.expect_call(timeout=None).and_return(1)
    295             (subcommand.cPickle.load.expect_call(task.result_pickle)
    296                     .and_return(6))
    297             task.result_pickle.close.expect_call()
    298 
    299         self.assertRaises(subcommand.error.AutoservError, subcommand.parallel,
    300                           tasklist)
    301         self.god.check_playback()
    302 
    303 
    304     def test_timeout(self):
    305         self.god.stub_function(subcommand.time, 'time')
    306 
    307         tasklist = self._setup_common()
    308         timeout = 10
    309 
    310         subcommand.time.time.expect_call().and_return(1)
    311 
    312         for task in tasklist:
    313             subcommand.time.time.expect_call().and_return(1)
    314             task.fork_waitfor.expect_call(timeout=timeout).and_return(None)
    315             (subcommand.cPickle.load.expect_call(task.result_pickle)
    316                     .and_return(6))
    317             task.result_pickle.close.expect_call()
    318 
    319         self.assertRaises(subcommand.error.AutoservError, subcommand.parallel,
    320                           tasklist, timeout=timeout)
    321         self.god.check_playback()
    322 
    323 
    324     def test_return_results(self):
    325         tasklist = self._setup_common()
    326 
    327         tasklist[0].fork_waitfor.expect_call(timeout=None).and_return(0)
    328         (subcommand.cPickle.load.expect_call(tasklist[0].result_pickle)
    329                 .and_return(6))
    330         tasklist[0].result_pickle.close.expect_call()
    331 
    332         error = Exception('fail')
    333         tasklist[1].fork_waitfor.expect_call(timeout=None).and_return(1)
    334         (subcommand.cPickle.load.expect_call(tasklist[1].result_pickle)
    335                 .and_return(error))
    336         tasklist[1].result_pickle.close.expect_call()
    337 
    338         self.assertEquals(subcommand.parallel(tasklist, return_results=True),
    339                           [6, error])
    340         self.god.check_playback()
    341 
    342 
    343 class test_parallel_simple(unittest.TestCase):
    344     def setUp(self):
    345         self.god = mock.mock_god()
    346         self.god.stub_function(subcommand, 'parallel')
    347         ctor = self.god.create_mock_function('subcommand')
    348         self.god.stub_with(subcommand, 'subcommand', ctor)
    349 
    350 
    351     def tearDown(self):
    352         self.god.unstub_all()
    353 
    354 
    355     def test_simple_success(self):
    356         func = self.god.create_mock_function('func')
    357 
    358         func.expect_call(3)
    359 
    360         subcommand.parallel_simple(func, (3,))
    361         self.god.check_playback()
    362 
    363 
    364     def test_simple_failure(self):
    365         func = self.god.create_mock_function('func')
    366 
    367         error = Exception('fail')
    368         func.expect_call(3).and_raises(error)
    369 
    370         self.assertRaises(Exception, subcommand.parallel_simple, func, (3,))
    371         self.god.check_playback()
    372 
    373 
    374     def test_simple_return_value(self):
    375         func = self.god.create_mock_function('func')
    376 
    377         result = 1000
    378         func.expect_call(3).and_return(result)
    379 
    380         self.assertEquals(subcommand.parallel_simple(func, (3,),
    381                                                      return_results=True),
    382                           [result])
    383         self.god.check_playback()
    384 
    385 
    386     def test_default_subdirs_constructor(self):
    387         func = self.god.create_mock_function('func')
    388         args = range(4)
    389         for arg in args:
    390             subcommand.subcommand.expect_call(
    391                     func, [arg], str(arg)).and_return(arg)
    392         subcommand.parallel.expect_call(args, None, return_results=False)
    393 
    394         subcommand.parallel_simple(func, args)
    395         self.god.check_playback()
    396 
    397 
    398     def test_nolog_skips_subdirs(self):
    399         func = self.god.create_mock_function('func')
    400         args = range(3)
    401         for arg in args:
    402             subcommand.subcommand.expect_call(
    403                     func, [arg], None).and_return(arg)
    404         subcommand.parallel.expect_call(args, None, return_results=False)
    405 
    406         subcommand.parallel_simple(func, args, log=False)
    407         self.god.check_playback()
    408 
    409 
    410     def test_custom_subdirs_constructor(self):
    411         func = self.god.create_mock_function('func')
    412         args = range(7)
    413         subdirs = ['subdir%s' % arg for arg in args]
    414         for arg, subdir in zip(args, subdirs):
    415             subcommand.subcommand.expect_call(
    416                     func, [arg], subdir).and_return(arg)
    417         subcommand.parallel.expect_call(args, None, return_results=False)
    418 
    419         subcommand.parallel_simple(
    420                 func, args, subdir_name_constructor=lambda x: 'subdir%s' % x)
    421         self.god.check_playback()
    422 
    423 
    424 if __name__ == '__main__':
    425     unittest.main()
    426