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