• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python
2# Copyright 2009 Google Inc. Released under the GPL v2
3
4import time, unittest
5
6import common
7from autotest_lib.client.common_lib import error
8from autotest_lib.client.common_lib.test_utils import mock
9from autotest_lib.server import subcommand
10
11
12def _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
26class 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
211class 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
247class 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
343class 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
424if __name__ == '__main__':
425    unittest.main()
426