1#!/usr/bin/python 2# Copyright 2016 The Chromium OS 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 mock 7import unittest 8 9import common 10 11from autotest_lib.server.cros.dynamic_suite import frontend_wrappers 12from autotest_lib.server import utils 13from autotest_lib.server.hosts import base_label 14 15 # pylint: disable=missing-docstring 16 17 18class TestBaseLabel(base_label.BaseLabel): 19 """TestBaseLabel is used for testing/validating BaseLabel methods.""" 20 21 _NAME = 'base_label' 22 23 def exists(self, host): 24 return host.exists 25 26 27class TestBaseLabels(base_label.BaseLabel): 28 """ 29 TestBaseLabels is used for testing/validating BaseLabel methods. 30 31 This is a variation of BaseLabel with multiple labels for _NAME 32 to ensure we handle a label that contains a list of labels for 33 its _NAME attribute. 34 """ 35 36 _NAME = ['base_label_1' , 'base_label_2'] 37 38 39class TestStringPrefixLabel(base_label.StringPrefixLabel): 40 """ 41 TestBaseLabels is used for testing/validating StringPrefixLabel methods. 42 43 This test class is to check that we properly construct the prefix labels 44 with the label passed in during construction. 45 """ 46 47 _NAME = 'prefix' 48 49 def __init__(self, label='postfix'): 50 self.label_to_return = label 51 52 53 def generate_labels(self, _): 54 return [self.label_to_return] 55 56 57class TestLabelList(base_label.StringLabel): 58 """TestLabelList is used to validate a label consisting of other labels.""" 59 60 _LABEL_LIST = [TestBaseLabel(), TestStringPrefixLabel()] 61 62 def generate_labels(self, host): 63 labels = [] 64 for label_detectors in self._LABEL_LIST: 65 labels.extend(label_detectors.get(host)) 66 return labels 67 68 69class MockAFEHost(utils.EmptyAFEHost): 70 71 def __init__(self, labels=[], attributes={}): 72 self.labels = labels 73 self.attributes = attributes 74 75 76class MockHost(object): 77 78 def __init__(self, exists=True, afe_host=None): 79 self.hostname = 'hostname' 80 self.exists = exists 81 self._afe_host = afe_host 82 83 84class BaseLabelUnittests(unittest.TestCase): 85 """Unittest for testing base_label.BaseLabel.""" 86 87 def setUp(self): 88 self.test_base_label = TestBaseLabel() 89 self.test_base_labels = TestBaseLabels() 90 91 92 def test_generate_labels(self): 93 """Let's make sure generate_labels() returns the labels expected.""" 94 self.assertEqual(self.test_base_label.generate_labels(None), 95 [self.test_base_label._NAME]) 96 97 98 def test_get(self): 99 """Let's make sure the logic in get() works as expected.""" 100 # We should get labels here. 101 self.assertEqual(self.test_base_label.get(MockHost(exists=True)), 102 [self.test_base_label._NAME]) 103 # We should get nothing here. 104 self.assertEqual(self.test_base_label.get(MockHost(exists=False)), 105 []) 106 107 108 def test_get_all_labels(self): 109 """Check that we get the expected labels for get_all_labels().""" 110 prefix_tbl, full_tbl = self.test_base_label.get_all_labels() 111 prefix_tbls, full_tbls = self.test_base_labels.get_all_labels() 112 113 # We want to check that we always get a list of labels regardless if 114 # the label class attribute _NAME is a list or a string. 115 self.assertEqual(full_tbl, set([self.test_base_label._NAME])) 116 self.assertEqual(full_tbls, set(self.test_base_labels._NAME)) 117 118 # We want to make sure we get nothing on the prefix_* side of things 119 # since BaseLabel shouldn't be a prefix for any label. 120 self.assertEqual(prefix_tbl, set()) 121 self.assertEqual(prefix_tbls, set()) 122 123 124class StringPrefixLabelUnittests(unittest.TestCase): 125 """Unittest for testing base_label.StringPrefixLabel.""" 126 127 def setUp(self): 128 self.postfix_label = 'postfix_label' 129 self.test_label = TestStringPrefixLabel(label=self.postfix_label) 130 131 132 def test_get(self): 133 """Let's make sure that the labels we get are prefixed.""" 134 self.assertEqual(self.test_label.get(None), 135 ['%s:%s' % (self.test_label._NAME, 136 self.postfix_label)]) 137 138 139 def test_get_all_labels(self): 140 """Check that we only get prefix labels and no full labels.""" 141 prefix_labels, postfix_labels = self.test_label.get_all_labels() 142 self.assertEqual(prefix_labels, set(['%s:' % self.test_label._NAME])) 143 self.assertEqual(postfix_labels, set()) 144 145 146class LabelRetrieverUnittests(unittest.TestCase): 147 """Unittest for testing base_label.LabelRetriever.""" 148 149 def setUp(self): 150 label_list = [TestStringPrefixLabel(), TestBaseLabel()] 151 self.retriever = base_label.LabelRetriever(label_list) 152 self.retriever._populate_known_labels(label_list) 153 self.retriever_label_list = base_label.LabelRetriever([TestLabelList()]) 154 self.retriever_label_list._populate_known_labels([TestLabelList()]) 155 156 157 def test_populate_known_labels(self): 158 """Check that _populate_known_labels() works as expected.""" 159 full_names = set([TestBaseLabel._NAME]) 160 prefix_names = set(['%s:' % TestStringPrefixLabel._NAME]) 161 # Check on a normal retriever. 162 self.assertEqual(self.retriever.label_full_names, full_names) 163 self.assertEqual(self.retriever.label_prefix_names, prefix_names) 164 165 # Check on a retriever that has a label with a label list. 166 self.assertEqual(self.retriever_label_list.label_full_names, full_names) 167 self.assertEqual(self.retriever_label_list.label_prefix_names, 168 prefix_names) 169 170 171 def test_is_known_label(self): 172 """Check _is_known_label() detects/skips the right labels.""" 173 # This will be a list of tuples of label and expected return bool. 174 # Make sure Full matches match correctly 175 labels_to_check = [(TestBaseLabel._NAME, True), 176 ('%s:' % TestStringPrefixLabel._NAME, True), 177 # Make sure partial matches fail. 178 (TestBaseLabel._NAME[:2], False), 179 ('%s:' % TestStringPrefixLabel._NAME[:2], False), 180 ('no_label_match', False)] 181 182 for label, expected_known in labels_to_check: 183 self.assertEqual(self.retriever._is_known_label(label), 184 expected_known) 185 186 187 @mock.patch.object(frontend_wrappers, 'RetryingAFE') 188 def test_update_labels(self, mock_retry_afe): 189 """Check that we add/remove the expected labels in update_labels().""" 190 label_to_add = 'label_to_add' 191 label_to_remove = 'prefix:label_to_remove' 192 mock_afe = mock.MagicMock() 193 mockhost = MockHost(afe_host=MockAFEHost( 194 labels=[label_to_remove, TestBaseLabel._NAME])) 195 expected_remove_labels = [label_to_remove] 196 expected_add_labels = ['%s:%s' % (TestStringPrefixLabel._NAME, 197 label_to_add)] 198 199 retriever = base_label.LabelRetriever( 200 [TestStringPrefixLabel(label=label_to_add), 201 TestBaseLabel()]) 202 203 retriever.update_labels(mockhost) 204 205 # Check that we removed the right labels 206 mock_afe.run.has_calls('host_remove_labels', 207 id=mockhost.hostname, 208 labels=expected_remove_labels) 209 210 # Check that we added the right labels 211 mock_afe.run.has_calls('host_add_labels', 212 id=mockhost.hostname, 213 labels=expected_add_labels) 214 215 216if __name__ == '__main__': 217 unittest.main() 218 219