mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 1 | #!/usr/bin/python |
| 2 | # Copyright 2009 Google Inc. Released under the GPL v2 |
| 3 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 4 | import time, unittest |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 5 | |
| 6 | import common |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 7 | from autotest_lib.client.common_lib import error |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 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 |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 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 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 211 | class real_subcommand_test(unittest.TestCase): |
| 212 | """Test actually running subcommands (without mocking).""" |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 213 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 214 | |
| 215 | def _setup_subcommand(self, func, *args): |
| 216 | cmd = subcommand.subcommand(func, args) |
| 217 | cmd.fork_start() |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 218 | return cmd |
| 219 | |
| 220 | |
| 221 | def test_fork_waitfor_no_timeout(self): |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 222 | """Test fork_waitfor success with no timeout.""" |
| 223 | cmd = self._setup_subcommand(lambda: None) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 224 | self.assertEquals(cmd.fork_waitfor(), 0) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 225 | |
| 226 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 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) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 231 | |
| 232 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 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) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 238 | |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 239 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 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) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 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 | |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 386 | def test_default_subdirs_constructor(self): |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 387 | func = self.god.create_mock_function('func') |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 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) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 393 | |
| 394 | subcommand.parallel_simple(func, args) |
| 395 | self.god.check_playback() |
| 396 | |
| 397 | |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 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) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 405 | |
| 406 | subcommand.parallel_simple(func, args, log=False) |
| 407 | self.god.check_playback() |
| 408 | |
| 409 | |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 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 | |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 424 | if __name__ == '__main__': |
| 425 | unittest.main() |