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