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