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 6 import mock 7 import unittest 8 9 import common 10 11 from autotest_lib.server import utils 12 from autotest_lib.server.hosts import host_info 13 from autotest_lib.server.hosts import base_label 14 15 # pylint: disable=missing-docstring 16 17 18 class 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 27 class 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 39 class 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 57 class MockAFEHost(utils.EmptyAFEHost): 58 59 def __init__(self, labels=None, attributes=None): 60 self.labels = labels or [] 61 self.attributes = attributes or {} 62 63 64 class MockHost(object): 65 66 def __init__(self, exists=True, store=None): 67 self.hostname = 'hostname' 68 self.exists = exists 69 self.host_info_store = store 70 71 72 class BaseLabelUnittests(unittest.TestCase): 73 """Unittest for testing base_label.BaseLabel.""" 74 75 def setUp(self): 76 self.test_base_label = TestBaseLabel() 77 self.test_base_labels = TestBaseLabels() 78 79 80 def test_generate_labels(self): 81 """Let's make sure generate_labels() returns the labels expected.""" 82 self.assertEqual(self.test_base_label.generate_labels(None), 83 [self.test_base_label._NAME]) 84 85 86 def test_get(self): 87 """Let's make sure the logic in get() works as expected.""" 88 # We should get labels here. 89 self.assertEqual(self.test_base_label.get(MockHost(exists=True)), 90 [self.test_base_label._NAME]) 91 # We should get nothing here. 92 self.assertEqual(self.test_base_label.get(MockHost(exists=False)), 93 []) 94 95 96 def test_get_all_labels(self): 97 """Check that we get the expected labels for get_all_labels().""" 98 prefix_tbl, full_tbl = self.test_base_label.get_all_labels() 99 prefix_tbls, full_tbls = self.test_base_labels.get_all_labels() 100 101 # We want to check that we always get a list of labels regardless if 102 # the label class attribute _NAME is a list or a string. 103 self.assertEqual(full_tbl, set([self.test_base_label._NAME])) 104 self.assertEqual(full_tbls, set(self.test_base_labels._NAME)) 105 106 # We want to make sure we get nothing on the prefix_* side of things 107 # since BaseLabel shouldn't be a prefix for any label. 108 self.assertEqual(prefix_tbl, set()) 109 self.assertEqual(prefix_tbls, set()) 110 111 112 class StringPrefixLabelUnittests(unittest.TestCase): 113 """Unittest for testing base_label.StringPrefixLabel.""" 114 115 def setUp(self): 116 self.postfix_label = 'postfix_label' 117 self.test_label = TestStringPrefixLabel(label=self.postfix_label) 118 119 120 def test_get(self): 121 """Let's make sure that the labels we get are prefixed.""" 122 self.assertEqual(self.test_label.get(None), 123 ['%s:%s' % (self.test_label._NAME, 124 self.postfix_label)]) 125 126 127 def test_get_all_labels(self): 128 """Check that we only get prefix labels and no full labels.""" 129 prefix_labels, postfix_labels = self.test_label.get_all_labels() 130 self.assertEqual(prefix_labels, set(['%s:' % self.test_label._NAME])) 131 self.assertEqual(postfix_labels, set()) 132 133 134 class LabelRetrieverUnittests(unittest.TestCase): 135 """Unittest for testing base_label.LabelRetriever.""" 136 137 def setUp(self): 138 label_list = [TestStringPrefixLabel(), TestBaseLabel()] 139 self.retriever = base_label.LabelRetriever(label_list) 140 self.retriever._populate_known_labels(label_list) 141 142 143 def test_populate_known_labels(self): 144 """Check that _populate_known_labels() works as expected.""" 145 full_names = set([TestBaseLabel._NAME]) 146 prefix_names = set(['%s:' % TestStringPrefixLabel._NAME]) 147 # Check on a normal retriever. 148 self.assertEqual(self.retriever.label_full_names, full_names) 149 self.assertEqual(self.retriever.label_prefix_names, prefix_names) 150 151 152 def test_is_known_label(self): 153 """Check _is_known_label() detects/skips the right labels.""" 154 # This will be a list of tuples of label and expected return bool. 155 # Make sure Full matches match correctly 156 labels_to_check = [(TestBaseLabel._NAME, True), 157 ('%s:' % TestStringPrefixLabel._NAME, True), 158 # Make sure partial matches fail. 159 (TestBaseLabel._NAME[:2], False), 160 ('%s:' % TestStringPrefixLabel._NAME[:2], False), 161 ('no_label_match', False)] 162 163 for label, expected_known in labels_to_check: 164 self.assertEqual(self.retriever._is_known_label(label), 165 expected_known) 166 167 168 def test_update_labels(self): 169 """Check that we add/remove the expected labels in update_labels().""" 170 label_to_add = 'label_to_add' 171 label_to_remove = 'prefix:label_to_remove' 172 store = host_info.InMemoryHostInfoStore( 173 info=host_info.HostInfo( 174 labels=[label_to_remove, TestBaseLabel._NAME], 175 ), 176 ) 177 mockhost = MockHost(store=store) 178 179 retriever = base_label.LabelRetriever( 180 [TestStringPrefixLabel(label=label_to_add), 181 TestBaseLabel()]) 182 retriever.update_labels(mockhost) 183 self.assertEqual( 184 set(store.get().labels), 185 {'%s:%s' % (TestStringPrefixLabel._NAME, label_to_add), 186 TestBaseLabel._NAME}, 187 ) 188 189 190 if __name__ == '__main__': 191 unittest.main() 192 193
