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