Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 1 | #!/usr/bin/env python |
| 2 | # Copyright 2016 The Chromium Authors. All rights reserved. |
| 3 | # Use of this source code is governed by a BSD-style license that can be |
| 4 | # found in the LICENSE file. |
| 5 | |
| 6 | import unittest |
| 7 | import mock |
Allen Li | de5ecce | 2017-11-06 17:48:09 -0800 | [diff] [blame] | 8 | |
Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 9 | import common |
Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 10 | from autotest_lib.site_utils import diagnosis_utils |
Allen Li | 035b47c | 2017-11-20 17:58:16 -0800 | [diff] [blame] | 11 | from autotest_lib.server import frontend |
Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 12 | |
| 13 | |
| 14 | class DiagnosisUtilsTest(unittest.TestCase): |
| 15 | """Tests for diagnosis_utils.""" |
| 16 | |
| 17 | def setUp(self): |
| 18 | """Set up test.""" |
| 19 | self.afe_mock = mock.MagicMock() |
Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 20 | |
| 21 | def _constructRPCHelper(self): |
| 22 | """Method to construct RPCHelper""" |
| 23 | return diagnosis_utils.RPCHelper(self.afe_mock) |
| 24 | |
| 25 | def _mockZeroHost(self): |
| 26 | """Mock method to return zero host""" |
| 27 | return () |
| 28 | |
| 29 | def _mockTwoAvailableHosts(self): |
| 30 | """Mock method to return two available hosts""" |
Allen Li | de5ecce | 2017-11-06 17:48:09 -0800 | [diff] [blame] | 31 | host_a = _StubHost(status='Ready', locked=False) |
| 32 | host_b = _StubHost(status='Ready', locked=False) |
Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 33 | return (host_a, host_b) |
| 34 | |
| 35 | def _mockTwoFailedHosts(self): |
| 36 | """Mock method to return two unavailable hosts""" |
Allen Li | de5ecce | 2017-11-06 17:48:09 -0800 | [diff] [blame] | 37 | host_a = _StubHost(status='Repair Failed', locked=False) |
| 38 | host_b = _StubHost(status='Repairing', locked=False) |
Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 39 | return (host_a, host_b) |
| 40 | |
| 41 | def testCheckDutAvailable(self): |
| 42 | """Test check_dut_availability with different scenarios""" |
| 43 | rpc_helper = self._constructRPCHelper() |
| 44 | board = 'test_board' |
| 45 | pool = 'test_pool' |
| 46 | |
| 47 | # Mock aef.get_hosts to return 0 host |
| 48 | self.afe_mock.get_hosts.return_value = self._mockZeroHost() |
| 49 | skip_duts_check = False |
| 50 | |
| 51 | # When minimum_duts is 0, do not force available DUTs |
| 52 | minimum_duts = 0 |
| 53 | rpc_helper.check_dut_availability(board, pool, |
| 54 | minimum_duts=minimum_duts, |
| 55 | skip_duts_check=skip_duts_check) |
| 56 | |
| 57 | # When skip_duts_check = False and minimum_duts > 0 and there's no host, |
| 58 | # should raise BoardNotAvailableError |
| 59 | minimum_duts = 1 |
| 60 | with self.assertRaises(diagnosis_utils.BoardNotAvailableError): |
| 61 | rpc_helper.check_dut_availability(board, pool, |
| 62 | minimum_duts=minimum_duts, |
| 63 | skip_duts_check=skip_duts_check) |
| 64 | |
| 65 | # Mock aef.get_hosts to return 2 available hosts |
| 66 | self.afe_mock.get_hosts.return_value = self._mockTwoAvailableHosts() |
| 67 | |
| 68 | # Set skip_duts_check to True, should not force checking avialble DUTs |
| 69 | # although available DUTs are less then minimum_duts |
| 70 | minimum_duts = 4 |
| 71 | skip_duts_check = True |
| 72 | rpc_helper.check_dut_availability(board, pool, |
| 73 | minimum_duts=minimum_duts, |
| 74 | skip_duts_check=skip_duts_check) |
| 75 | |
| 76 | # Set skip_duts_check to False and current available DUTs are less |
| 77 | # then minimum_duts, should raise NotEnoughDutsError |
| 78 | skip_duts_check = False |
| 79 | with self.assertRaises(diagnosis_utils.NotEnoughDutsError): |
| 80 | rpc_helper.check_dut_availability(board, pool, |
| 81 | minimum_duts=minimum_duts, |
| 82 | skip_duts_check=skip_duts_check) |
| 83 | |
| 84 | # When skip_duts_check is False and current available DUTs |
| 85 | # satisfy minimum_duts, no errors |
| 86 | minimum_duts = 2 |
| 87 | rpc_helper.check_dut_availability(board, pool, |
| 88 | minimum_duts=minimum_duts, |
| 89 | skip_duts_check=skip_duts_check) |
| 90 | |
| 91 | # Mock aef.get_hosts to return 2 failed hosts |
| 92 | self.afe_mock.get_hosts.return_value = self._mockTwoFailedHosts() |
| 93 | |
| 94 | # When skip_duts_check is False and the two hosts are not available, |
| 95 | # should raise NotEnoughDutsError |
| 96 | with self.assertRaises(diagnosis_utils.NotEnoughDutsError): |
| 97 | rpc_helper.check_dut_availability(board, pool, |
| 98 | minimum_duts=minimum_duts, |
| 99 | skip_duts_check=skip_duts_check) |
| 100 | |
Allen Li | de5ecce | 2017-11-06 17:48:09 -0800 | [diff] [blame] | 101 | |
| 102 | class _StubHost(object): |
| 103 | |
| 104 | def __init__(self, status, locked): |
| 105 | self.status = status |
| 106 | self.locked = locked |
| 107 | |
Allen Li | 035b47c | 2017-11-20 17:58:16 -0800 | [diff] [blame] | 108 | is_available = frontend.Host.is_available.im_func |
Allen Li | de5ecce | 2017-11-06 17:48:09 -0800 | [diff] [blame] | 109 | |
| 110 | |
Ningning Xia | f2c206c | 2016-04-13 14:15:51 -0700 | [diff] [blame] | 111 | if __name__ == '__main__': |
| 112 | unittest.main() |