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 8 9 import common 10 from autotest_lib.site_utils import diagnosis_utils 11 from autotest_lib.server import frontend 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() 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""" 31 host_a = _StubHost(status='Ready', locked=False) 32 host_b = _StubHost(status='Ready', locked=False) 33 return (host_a, host_b) 34 35 def _mockTwoFailedHosts(self): 36 """Mock method to return two unavailable hosts""" 37 host_a = _StubHost(status='Repair Failed', locked=False) 38 host_b = _StubHost(status='Repairing', locked=False) 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 101 102 class _StubHost(object): 103 104 def __init__(self, status, locked): 105 self.status = status 106 self.locked = locked 107 108 is_available = frontend.Host.is_available.im_func 109 110 111 if __name__ == '__main__': 112 unittest.main() 113
