xref: /aosp_15_r20/external/autotest/site_utils/diagnosis_utils_unittest.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
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