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            self.lambda_function = lambda: func(*args)
23
24    return wrapper(func, args)
25
26
27class 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
212class 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
248class 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
344class 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
425if __name__ == '__main__':
426    unittest.main()
427