433 lines
14 KiB
Python
433 lines
14 KiB
Python
|
#!/usr/bin/python2
|
||
|
|
# Copyright 2009 Google Inc. Released under the GPL v2
|
||
|
|
|
||
|
|
from __future__ import absolute_import
|
||
|
|
from __future__ import division
|
||
|
|
from __future__ import print_function
|
||
|
|
|
||
|
|
import time, unittest
|
||
|
|
import six
|
||
|
|
|
||
|
|
import common
|
||
|
|
from autotest_lib.client.common_lib import error
|
||
|
|
from autotest_lib.client.common_lib.test_utils import mock
|
||
|
|
from autotest_lib.server import subcommand
|
||
|
|
from six.moves import range
|
||
|
|
from six.moves import zip
|
||
|
|
|
||
|
|
|
||
|
|
def _create_subcommand(func, args):
|
||
|
|
# to avoid __init__
|
||
|
|
class wrapper(subcommand.subcommand):
|
||
|
|
def __init__(self, func, args):
|
||
|
|
self.func = func
|
||
|
|
self.args = args
|
||
|
|
self.subdir = None
|
||
|
|
self.debug = None
|
||
|
|
self.pid = None
|
||
|
|
self.returncode = None
|
||
|
|
|
||
|
|
return wrapper(func, args)
|
||
|
|
|
||
|
|
|
||
|
|
class subcommand_test(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.god = mock.mock_god()
|
||
|
|
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.god.unstub_all()
|
||
|
|
# cleanup the hooks
|
||
|
|
subcommand.subcommand.fork_hooks = []
|
||
|
|
subcommand.subcommand.join_hooks = []
|
||
|
|
|
||
|
|
|
||
|
|
def test_create(self):
|
||
|
|
def check_attributes(cmd, func, args, subdir=None, debug=None,
|
||
|
|
pid=None, returncode=None, fork_hooks=[],
|
||
|
|
join_hooks=[]):
|
||
|
|
self.assertEquals(cmd.func, func)
|
||
|
|
self.assertEquals(cmd.args, args)
|
||
|
|
self.assertEquals(cmd.subdir, subdir)
|
||
|
|
self.assertEquals(cmd.debug, debug)
|
||
|
|
self.assertEquals(cmd.pid, pid)
|
||
|
|
self.assertEquals(cmd.returncode, returncode)
|
||
|
|
self.assertEquals(cmd.fork_hooks, fork_hooks)
|
||
|
|
self.assertEquals(cmd.join_hooks, join_hooks)
|
||
|
|
|
||
|
|
def func(arg1, arg2):
|
||
|
|
pass
|
||
|
|
|
||
|
|
cmd = subcommand.subcommand(func, (2, 3))
|
||
|
|
check_attributes(cmd, func, (2, 3))
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
self.god.stub_function(subcommand.os.path, 'abspath')
|
||
|
|
self.god.stub_function(subcommand.os.path, 'exists')
|
||
|
|
self.god.stub_function(subcommand.os, 'mkdir')
|
||
|
|
|
||
|
|
subcommand.os.path.abspath.expect_call('dir').and_return('/foo/dir')
|
||
|
|
subcommand.os.path.exists.expect_call('/foo/dir').and_return(False)
|
||
|
|
subcommand.os.mkdir.expect_call('/foo/dir')
|
||
|
|
|
||
|
|
(subcommand.os.path.exists.expect_call('/foo/dir/debug')
|
||
|
|
.and_return(False))
|
||
|
|
subcommand.os.mkdir.expect_call('/foo/dir/debug')
|
||
|
|
|
||
|
|
cmd = subcommand.subcommand(func, (2, 3), subdir='dir')
|
||
|
|
check_attributes(cmd, func, (2, 3), subdir='/foo/dir',
|
||
|
|
debug='/foo/dir/debug')
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def _setup_fork_start_parent(self):
|
||
|
|
self.god.stub_function(subcommand.os, 'fork')
|
||
|
|
|
||
|
|
subcommand.os.fork.expect_call().and_return(1000)
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
cmd = _create_subcommand(func, [])
|
||
|
|
cmd.fork_start()
|
||
|
|
|
||
|
|
return cmd
|
||
|
|
|
||
|
|
|
||
|
|
def test_fork_start_parent(self):
|
||
|
|
cmd = self._setup_fork_start_parent()
|
||
|
|
|
||
|
|
self.assertEquals(cmd.pid, 1000)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def _setup_fork_start_child(self):
|
||
|
|
self.god.stub_function(subcommand.os, 'pipe')
|
||
|
|
self.god.stub_function(subcommand.os, 'fork')
|
||
|
|
self.god.stub_function(subcommand.os, 'close')
|
||
|
|
self.god.stub_function(subcommand.os, 'write')
|
||
|
|
self.god.stub_function(six.moves.cPickle, 'dumps')
|
||
|
|
self.god.stub_function(subcommand.os, '_exit')
|
||
|
|
|
||
|
|
|
||
|
|
def test_fork_start_child(self):
|
||
|
|
self._setup_fork_start_child()
|
||
|
|
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
fork_hook = self.god.create_mock_function('fork_hook')
|
||
|
|
join_hook = self.god.create_mock_function('join_hook')
|
||
|
|
|
||
|
|
subcommand.subcommand.register_fork_hook(fork_hook)
|
||
|
|
subcommand.subcommand.register_join_hook(join_hook)
|
||
|
|
cmd = _create_subcommand(func, (1, 2))
|
||
|
|
|
||
|
|
subcommand.os.pipe.expect_call().and_return((10, 20))
|
||
|
|
subcommand.os.fork.expect_call().and_return(0)
|
||
|
|
subcommand.os.close.expect_call(10)
|
||
|
|
fork_hook.expect_call(cmd)
|
||
|
|
func.expect_call(1, 2).and_return(True)
|
||
|
|
six.moves.cPickle.dumps.expect_call(True,
|
||
|
|
six.moves.cPickle.HIGHEST_PROTOCOL).and_return('True')
|
||
|
|
subcommand.os.write.expect_call(20, 'True')
|
||
|
|
subcommand.os.close.expect_call(20)
|
||
|
|
join_hook.expect_call(cmd)
|
||
|
|
subcommand.os._exit.expect_call(0)
|
||
|
|
|
||
|
|
cmd.fork_start()
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_fork_start_child_error(self):
|
||
|
|
self._setup_fork_start_child()
|
||
|
|
self.god.stub_function(subcommand.logging, 'exception')
|
||
|
|
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
cmd = _create_subcommand(func, (1, 2))
|
||
|
|
error = Exception('some error')
|
||
|
|
|
||
|
|
subcommand.os.pipe.expect_call().and_return((10, 20))
|
||
|
|
subcommand.os.fork.expect_call().and_return(0)
|
||
|
|
subcommand.os.close.expect_call(10)
|
||
|
|
func.expect_call(1, 2).and_raises(error)
|
||
|
|
subcommand.logging.exception.expect_call('function failed')
|
||
|
|
six.moves.cPickle.dumps.expect_call(error,
|
||
|
|
six.moves.cPickle.HIGHEST_PROTOCOL).and_return('error')
|
||
|
|
subcommand.os.write.expect_call(20, 'error')
|
||
|
|
subcommand.os.close.expect_call(20)
|
||
|
|
subcommand.os._exit.expect_call(1)
|
||
|
|
|
||
|
|
cmd.fork_start()
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def _setup_poll(self):
|
||
|
|
cmd = self._setup_fork_start_parent()
|
||
|
|
self.god.stub_function(subcommand.os, 'waitpid')
|
||
|
|
return cmd
|
||
|
|
|
||
|
|
|
||
|
|
def test_poll_running(self):
|
||
|
|
cmd = self._setup_poll()
|
||
|
|
|
||
|
|
(subcommand.os.waitpid.expect_call(1000, subcommand.os.WNOHANG)
|
||
|
|
.and_raises(subcommand.os.error('waitpid')))
|
||
|
|
self.assertEquals(cmd.poll(), None)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_poll_finished_success(self):
|
||
|
|
cmd = self._setup_poll()
|
||
|
|
|
||
|
|
(subcommand.os.waitpid.expect_call(1000, subcommand.os.WNOHANG)
|
||
|
|
.and_return((1000, 0)))
|
||
|
|
self.assertEquals(cmd.poll(), 0)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_poll_finished_failure(self):
|
||
|
|
cmd = self._setup_poll()
|
||
|
|
self.god.stub_function(cmd, '_handle_exitstatus')
|
||
|
|
|
||
|
|
(subcommand.os.waitpid.expect_call(1000, subcommand.os.WNOHANG)
|
||
|
|
.and_return((1000, 10)))
|
||
|
|
cmd._handle_exitstatus.expect_call(10).and_raises(Exception('fail'))
|
||
|
|
|
||
|
|
self.assertRaises(Exception, cmd.poll)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_wait_success(self):
|
||
|
|
cmd = self._setup_poll()
|
||
|
|
|
||
|
|
(subcommand.os.waitpid.expect_call(1000, 0)
|
||
|
|
.and_return((1000, 0)))
|
||
|
|
|
||
|
|
self.assertEquals(cmd.wait(), 0)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_wait_failure(self):
|
||
|
|
cmd = self._setup_poll()
|
||
|
|
self.god.stub_function(cmd, '_handle_exitstatus')
|
||
|
|
|
||
|
|
(subcommand.os.waitpid.expect_call(1000, 0)
|
||
|
|
.and_return((1000, 10)))
|
||
|
|
|
||
|
|
cmd._handle_exitstatus.expect_call(10).and_raises(Exception('fail'))
|
||
|
|
self.assertRaises(Exception, cmd.wait)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
class real_subcommand_test(unittest.TestCase):
|
||
|
|
"""Test actually running subcommands (without mocking)."""
|
||
|
|
|
||
|
|
|
||
|
|
def _setup_subcommand(self, func, *args):
|
||
|
|
cmd = subcommand.subcommand(func, args)
|
||
|
|
cmd.fork_start()
|
||
|
|
return cmd
|
||
|
|
|
||
|
|
|
||
|
|
def test_fork_waitfor_no_timeout(self):
|
||
|
|
"""Test fork_waitfor success with no timeout."""
|
||
|
|
cmd = self._setup_subcommand(lambda: None)
|
||
|
|
self.assertEquals(cmd.fork_waitfor(), 0)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fork_waitfor_timeout(self):
|
||
|
|
"""Test fork_waitfor success with a timeout."""
|
||
|
|
cmd = self._setup_subcommand(lambda: None)
|
||
|
|
self.assertEquals(cmd.fork_waitfor(timeout=60), 0)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fork_waitfor_exception(self):
|
||
|
|
"""Test fork_waitfor failure with an exception."""
|
||
|
|
cmd = self._setup_subcommand(lambda: None, 'foo')
|
||
|
|
with self.assertRaises(error.AutoservSubcommandError):
|
||
|
|
cmd.fork_waitfor(timeout=60)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fork_waitfor_timeout_fail(self):
|
||
|
|
"""Test fork_waitfor timing out."""
|
||
|
|
cmd = self._setup_subcommand(lambda: time.sleep(60))
|
||
|
|
with self.assertRaises(error.AutoservSubcommandError):
|
||
|
|
cmd.fork_waitfor(timeout=1)
|
||
|
|
|
||
|
|
|
||
|
|
class parallel_test(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.god = mock.mock_god()
|
||
|
|
self.god.stub_function(six.moves.cPickle, 'load')
|
||
|
|
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.god.unstub_all()
|
||
|
|
|
||
|
|
|
||
|
|
def _get_cmd(self, func, args):
|
||
|
|
cmd = _create_subcommand(func, args)
|
||
|
|
cmd.result_pickle = self.god.create_mock_class(file, 'file')
|
||
|
|
return self.god.create_mock_class(cmd, 'subcommand')
|
||
|
|
|
||
|
|
|
||
|
|
def _get_tasklist(self):
|
||
|
|
return [self._get_cmd(lambda x: x * 2, (3,)),
|
||
|
|
self._get_cmd(lambda: None, [])]
|
||
|
|
|
||
|
|
|
||
|
|
def _setup_common(self):
|
||
|
|
tasklist = self._get_tasklist()
|
||
|
|
|
||
|
|
for task in tasklist:
|
||
|
|
task.fork_start.expect_call()
|
||
|
|
|
||
|
|
return tasklist
|
||
|
|
|
||
|
|
|
||
|
|
def test_success(self):
|
||
|
|
tasklist = self._setup_common()
|
||
|
|
|
||
|
|
for task in tasklist:
|
||
|
|
task.fork_waitfor.expect_call(timeout=None).and_return(0)
|
||
|
|
(six.moves.cPickle.load.expect_call(task.result_pickle)
|
||
|
|
.and_return(6))
|
||
|
|
task.result_pickle.close.expect_call()
|
||
|
|
|
||
|
|
subcommand.parallel(tasklist)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_failure(self):
|
||
|
|
tasklist = self._setup_common()
|
||
|
|
|
||
|
|
for task in tasklist:
|
||
|
|
task.fork_waitfor.expect_call(timeout=None).and_return(1)
|
||
|
|
(six.moves.cPickle.load.expect_call(task.result_pickle)
|
||
|
|
.and_return(6))
|
||
|
|
task.result_pickle.close.expect_call()
|
||
|
|
|
||
|
|
self.assertRaises(subcommand.error.AutoservError, subcommand.parallel,
|
||
|
|
tasklist)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_timeout(self):
|
||
|
|
self.god.stub_function(subcommand.time, 'time')
|
||
|
|
|
||
|
|
tasklist = self._setup_common()
|
||
|
|
timeout = 10
|
||
|
|
|
||
|
|
subcommand.time.time.expect_call().and_return(1)
|
||
|
|
|
||
|
|
for task in tasklist:
|
||
|
|
subcommand.time.time.expect_call().and_return(1)
|
||
|
|
task.fork_waitfor.expect_call(timeout=timeout).and_return(None)
|
||
|
|
(six.moves.cPickle.load.expect_call(task.result_pickle)
|
||
|
|
.and_return(6))
|
||
|
|
task.result_pickle.close.expect_call()
|
||
|
|
|
||
|
|
self.assertRaises(subcommand.error.AutoservError, subcommand.parallel,
|
||
|
|
tasklist, timeout=timeout)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_return_results(self):
|
||
|
|
tasklist = self._setup_common()
|
||
|
|
|
||
|
|
tasklist[0].fork_waitfor.expect_call(timeout=None).and_return(0)
|
||
|
|
(six.moves.cPickle.load.expect_call(tasklist[0].result_pickle)
|
||
|
|
.and_return(6))
|
||
|
|
tasklist[0].result_pickle.close.expect_call()
|
||
|
|
|
||
|
|
error = Exception('fail')
|
||
|
|
tasklist[1].fork_waitfor.expect_call(timeout=None).and_return(1)
|
||
|
|
(six.moves.cPickle.load.expect_call(tasklist[1].result_pickle)
|
||
|
|
.and_return(error))
|
||
|
|
tasklist[1].result_pickle.close.expect_call()
|
||
|
|
|
||
|
|
self.assertEquals(subcommand.parallel(tasklist, return_results=True),
|
||
|
|
[6, error])
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
class test_parallel_simple(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.god = mock.mock_god()
|
||
|
|
self.god.stub_function(subcommand, 'parallel')
|
||
|
|
ctor = self.god.create_mock_function('subcommand')
|
||
|
|
self.god.stub_with(subcommand, 'subcommand', ctor)
|
||
|
|
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.god.unstub_all()
|
||
|
|
|
||
|
|
|
||
|
|
def test_simple_success(self):
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
|
||
|
|
func.expect_call(3)
|
||
|
|
|
||
|
|
subcommand.parallel_simple(func, (3,))
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_simple_failure(self):
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
|
||
|
|
error = Exception('fail')
|
||
|
|
func.expect_call(3).and_raises(error)
|
||
|
|
|
||
|
|
self.assertRaises(Exception, subcommand.parallel_simple, func, (3,))
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_simple_return_value(self):
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
|
||
|
|
result = 1000
|
||
|
|
func.expect_call(3).and_return(result)
|
||
|
|
|
||
|
|
self.assertEquals(subcommand.parallel_simple(func, (3,),
|
||
|
|
return_results=True),
|
||
|
|
[result])
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_default_subdirs_constructor(self):
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
args = list(range(4))
|
||
|
|
for arg in args:
|
||
|
|
subcommand.subcommand.expect_call(
|
||
|
|
func, [arg], str(arg)).and_return(arg)
|
||
|
|
subcommand.parallel.expect_call(args, None, return_results=False)
|
||
|
|
|
||
|
|
subcommand.parallel_simple(func, args)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_nolog_skips_subdirs(self):
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
args = list(range(3))
|
||
|
|
for arg in args:
|
||
|
|
subcommand.subcommand.expect_call(
|
||
|
|
func, [arg], None).and_return(arg)
|
||
|
|
subcommand.parallel.expect_call(args, None, return_results=False)
|
||
|
|
|
||
|
|
subcommand.parallel_simple(func, args, log=False)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
def test_custom_subdirs_constructor(self):
|
||
|
|
func = self.god.create_mock_function('func')
|
||
|
|
args = list(range(7))
|
||
|
|
subdirs = ['subdir%s' % arg for arg in args]
|
||
|
|
for arg, subdir in zip(args, subdirs):
|
||
|
|
subcommand.subcommand.expect_call(
|
||
|
|
func, [arg], subdir).and_return(arg)
|
||
|
|
subcommand.parallel.expect_call(args, None, return_results=False)
|
||
|
|
|
||
|
|
subcommand.parallel_simple(
|
||
|
|
func, args, subdir_name_constructor=lambda x: 'subdir%s' % x)
|
||
|
|
self.god.check_playback()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|