1# Copyright 2017 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import mock 6import unittest 7 8import common 9from autotest_lib.frontend.afe.json_rpc import proxy as rpc_proxy 10from autotest_lib.server import frontend 11from autotest_lib.server.hosts import afe_store 12from autotest_lib.server.hosts import host_info 13 14class AfeStoreTest(unittest.TestCase): 15 """Test refresh/commit success cases for AfeStore.""" 16 17 def setUp(self): 18 self.hostname = 'some-host' 19 self.mock_afe = mock.create_autospec(frontend.AFE, instance=True) 20 self.store = afe_store.AfeStore(self.hostname, afe=self.mock_afe) 21 22 23 def _create_mock_host(self, labels, attributes): 24 """Create a mock frontend.Host with the given labels and attributes. 25 26 @param labels: The labels to set on the host. 27 @param attributes: The attributes to set on the host. 28 @returns: A mock object for frontend.Host. 29 """ 30 mock_host = mock.create_autospec(frontend.Host, instance=True) 31 mock_host.labels = labels 32 mock_host.attributes = attributes 33 return mock_host 34 35 36 def test_refresh(self): 37 """Test that refresh correctly translates host information.""" 38 self.mock_afe.get_hosts.return_value = [ 39 self._create_mock_host(['label1'], {'attrib1': 'val1'})] 40 info = self.store._refresh_impl() 41 self.assertListEqual(info.labels, ['label1']) 42 self.assertDictEqual(info.attributes, {'attrib1': 'val1'}) 43 44 45 def test_refresh_no_host_raises(self): 46 """Test that refresh complains if no host is found.""" 47 self.mock_afe.get_hosts.return_value = [] 48 with self.assertRaises(host_info.StoreError): 49 self.store._refresh_impl() 50 51 52 def test_refresh_multiple_hosts_picks_first(self): 53 """Test that refresh returns the first host if multiple match.""" 54 self.mock_afe.get_hosts.return_value = [ 55 self._create_mock_host(['label1'], {'attrib1': 'val1'}), 56 self._create_mock_host(['label2'], {'attrib2': 'val2'})] 57 info = self.store._refresh_impl() 58 self.assertListEqual(info.labels, ['label1']) 59 self.assertDictEqual(info.attributes, {'attrib1': 'val1'}) 60 61 62 def test_commit_labels(self): 63 """Tests that labels are updated correctly on commit.""" 64 self.mock_afe.get_hosts.return_value = [ 65 self._create_mock_host(['label1'], {})] 66 info = host_info.HostInfo(['label2'], {}) 67 self.store._commit_impl(info) 68 self.assertEqual(self.mock_afe.run.call_count, 2) 69 expected_run_calls = [ 70 mock.call('host_remove_labels', id='some-host', 71 labels=['label1']), 72 mock.call('host_add_labels', id='some-host', 73 labels=['label2']), 74 ] 75 self.mock_afe.run.assert_has_calls(expected_run_calls, 76 any_order=True) 77 78 79 def test_commit_labels_raises(self): 80 """Test that exception while committing is translated properly.""" 81 self.mock_afe.get_hosts.return_value = [ 82 self._create_mock_host(['label1'], {})] 83 self.mock_afe.run.side_effect = rpc_proxy.JSONRPCException('some error') 84 info = host_info.HostInfo(['label2'], {}) 85 with self.assertRaises(host_info.StoreError): 86 self.store._commit_impl(info) 87 88 89 def test_commit_adds_attributes(self): 90 """Tests that new attributes are added correctly on commit.""" 91 self.mock_afe.get_hosts.return_value = [ 92 self._create_mock_host([], {})] 93 info = host_info.HostInfo([], {'attrib1': 'val1'}) 94 self.store._commit_impl(info) 95 self.assertEqual(self.mock_afe.set_host_attribute.call_count, 1) 96 self.mock_afe.set_host_attribute.assert_called_once_with( 97 'attrib1', 'val1', hostname=self.hostname) 98 99 100 def test_commit_updates_attributes(self): 101 """Tests that existing attributes are updated correctly on commit.""" 102 self.mock_afe.get_hosts.return_value = [ 103 self._create_mock_host([], {'attrib1': 'val1'})] 104 info = host_info.HostInfo([], {'attrib1': 'val1_updated'}) 105 self.store._commit_impl(info) 106 self.assertEqual(self.mock_afe.set_host_attribute.call_count, 1) 107 self.mock_afe.set_host_attribute.assert_called_once_with( 108 'attrib1', 'val1_updated', hostname=self.hostname) 109 110 111 def test_commit_deletes_attributes(self): 112 """Tests that deleted attributes are updated correctly on commit.""" 113 self.mock_afe.get_hosts.return_value = [ 114 self._create_mock_host([], {'attrib1': 'val1'})] 115 info = host_info.HostInfo([], {}) 116 self.store._commit_impl(info) 117 self.assertEqual(self.mock_afe.set_host_attribute.call_count, 1) 118 self.mock_afe.set_host_attribute.assert_called_once_with( 119 'attrib1', None, hostname=self.hostname) 120 121 122 def test_str(self): 123 """Sanity tests the __str__ implementaiton""" 124 self.assertEqual(str(self.store), 'AfeStore[some-host]') 125 126 127class AfeStoreKeepPoolTest(unittest.TestCase): 128 """Test commit success cases for AfeStoreKeepPool.""" 129 130 def setUp(self): 131 self.hostname = 'some-host' 132 self.mock_afe = mock.create_autospec(frontend.AFE, instance=True) 133 self.store = afe_store.AfeStoreKeepPool( 134 self.hostname, afe=self.mock_afe) 135 136 def _create_mock_host(self, labels, attributes): 137 """Create a mock frontend.Host with the given labels and attributes. 138 139 @param labels: The labels to set on the host. 140 @param attributes: The attributes to set on the host. 141 @returns: A mock object for frontend.Host. 142 """ 143 mock_host = mock.create_autospec(frontend.Host, instance=True) 144 mock_host.labels = labels 145 mock_host.attributes = attributes 146 return mock_host 147 148 def test_no_pool_label(self): 149 """Tests that no pool labels are updated on commit.""" 150 self.mock_afe.get_hosts.return_value = [ 151 self._create_mock_host(['label1'], {})] 152 new_info = host_info.HostInfo(['label2'], {}) 153 old_info = self.store._refresh_impl() 154 labels_to_remove, labels_to_add = self.store._adjust_pool( 155 old_info, new_info) 156 self.assertEqual((labels_to_remove, labels_to_add), 157 (['label1'], ['label2'])) 158 159 def test_update_pool_label(self): 160 """Tests that pool labels are updated correctly on commit.""" 161 self.mock_afe.get_hosts.return_value = [ 162 self._create_mock_host(['pool:XXX'], {})] 163 new_info = host_info.HostInfo(['pool:YYY'], {}) 164 old_info = self.store._refresh_impl() 165 labels_to_remove, labels_to_add = self.store._adjust_pool( 166 old_info, new_info) 167 self.assertEqual((labels_to_remove, labels_to_add), 168 (['pool:XXX'], ['pool:YYY'])) 169 170 def test_add_pool_label(self): 171 """Tests that pool labels are added correctly on commit.""" 172 self.mock_afe.get_hosts.return_value = [ 173 self._create_mock_host(['label1'], {})] 174 new_info = host_info.HostInfo(['pool:YYY'], {}) 175 old_info = self.store._refresh_impl() 176 labels_to_remove, labels_to_add = self.store._adjust_pool( 177 old_info, new_info) 178 self.assertEqual((labels_to_remove, labels_to_add), 179 (['label1'], ['pool:YYY'])) 180 181 def test_remove_pool_label(self): 182 """Tests that pool labels are not removed on commit.""" 183 self.mock_afe.get_hosts.return_value = [ 184 self._create_mock_host(['pool:XXX'], {})] 185 new_info = host_info.HostInfo(['label2'], {}) 186 old_info = self.store._refresh_impl() 187 labels_to_remove, labels_to_add = self.store._adjust_pool( 188 old_info, new_info) 189 self.assertEqual((labels_to_remove, labels_to_add), 190 ([], ['label2'])) 191 192 193class DictDiffTest(unittest.TestCase): 194 """Tests the afe_store._dict_diff private method.""" 195 196 def _assert_dict_diff(self, got_tuple, expectation_tuple): 197 """Verifies the result from _dict_diff 198 199 @param got_tuple: The tuple returned by afe_store._dict_diff 200 @param expectatin_tuple: tuple (left_only, right_only, differing) 201 containing iterable of keys to verify against got_tuple. 202 """ 203 for got, expect in zip(got_tuple, expectation_tuple): 204 self.assertEqual(got, set(expect)) 205 206 207 def test_both_empty(self): 208 """Tests the case when both dicts are empty.""" 209 self._assert_dict_diff(afe_store._dict_diff({}, {}), 210 ((), (), ())) 211 212 213 def test_right_dict_only(self): 214 """Tests the case when left dict is empty.""" 215 self._assert_dict_diff(afe_store._dict_diff({}, {1: 1}), 216 ((), (1,), ())) 217 218 219 def test_left_dict_only(self): 220 """Tests the case when right dict is empty.""" 221 self._assert_dict_diff(afe_store._dict_diff({1: 1}, {}), 222 ((1,), (), ())) 223 224 225 def test_left_dict_extra(self): 226 """Tests the case when left dict has extra keys.""" 227 self._assert_dict_diff(afe_store._dict_diff({1: 1, 2: 2}, {1: 1}), 228 ((2,), (), ())) 229 230 231 def test_right_dict_extra(self): 232 """Tests the case when right dict has extra keys.""" 233 self._assert_dict_diff(afe_store._dict_diff({1: 1}, {1: 1, 2: 2}), 234 ((), (2,), ())) 235 236 237 def test_identical_keys_with_different_values(self): 238 """Tests the case when the set of keys is same, but values differ.""" 239 self._assert_dict_diff(afe_store._dict_diff({1: 1, 2: 3}, {1: 1, 2: 2}), 240 ((), (), (2,))) 241 242 243if __name__ == '__main__': 244 unittest.main() 245