1#!/usr/bin/env python2 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 6import unittest 7import mock 8 9import common 10from autotest_lib.site_utils import diagnosis_utils 11from autotest_lib.server import frontend 12 13 14class 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 labels = ('board:test_board', 'pool:test_pool', 'model:test_model') 45 46 # Mock aef.get_hosts to return 0 host 47 self.afe_mock.get_hosts.return_value = self._mockZeroHost() 48 skip_duts_check = False 49 50 # When minimum_duts is 0, do not force available DUTs 51 minimum_duts = 0 52 rpc_helper.check_dut_availability(labels, 53 minimum_duts=minimum_duts, 54 skip_duts_check=skip_duts_check) 55 56 # When skip_duts_check = False and minimum_duts > 0 and there's no host, 57 # should raise DUTsNotAvailableError 58 minimum_duts = 1 59 with self.assertRaises(diagnosis_utils.DUTsNotAvailableError): 60 rpc_helper.check_dut_availability(labels, 61 minimum_duts=minimum_duts, 62 skip_duts_check=skip_duts_check) 63 64 # Mock aef.get_hosts to return 2 available hosts 65 self.afe_mock.get_hosts.return_value = self._mockTwoAvailableHosts() 66 67 # Set skip_duts_check to True, should not force checking avialble DUTs 68 # although available DUTs are less then minimum_duts 69 minimum_duts = 4 70 skip_duts_check = True 71 rpc_helper.check_dut_availability(labels, 72 minimum_duts=minimum_duts, 73 skip_duts_check=skip_duts_check) 74 75 # Set skip_duts_check to False and current available DUTs are less 76 # then minimum_duts, should raise NotEnoughDutsError 77 skip_duts_check = False 78 with self.assertRaises(diagnosis_utils.NotEnoughDutsError): 79 rpc_helper.check_dut_availability(labels, 80 minimum_duts=minimum_duts, 81 skip_duts_check=skip_duts_check) 82 83 # When skip_duts_check is False and current available DUTs 84 # satisfy minimum_duts, no errors 85 minimum_duts = 2 86 rpc_helper.check_dut_availability(labels, 87 minimum_duts=minimum_duts, 88 skip_duts_check=skip_duts_check) 89 90 # Mock aef.get_hosts to return 2 failed hosts 91 self.afe_mock.get_hosts.return_value = self._mockTwoFailedHosts() 92 93 # When skip_duts_check is False and the two hosts are not available, 94 # should raise NotEnoughDutsError 95 with self.assertRaises(diagnosis_utils.NotEnoughDutsError): 96 rpc_helper.check_dut_availability(labels, 97 minimum_duts=minimum_duts, 98 skip_duts_check=skip_duts_check) 99 100 101class _StubHost(object): 102 103 def __init__(self, status, locked): 104 self.status = status 105 self.locked = locked 106 107 is_available = frontend.Host.is_available.im_func 108 109 110if __name__ == '__main__': 111 unittest.main() 112