1#!/usr/bin/python 2# 3# Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7import mock 8import unittest 9 10import common 11from autotest_lib.frontend import setup_django_lite_environment 12from autotest_lib.frontend.afe import frontend_test_utils 13from autotest_lib.scheduler import rdb 14from autotest_lib.scheduler import rdb_hosts 15from autotest_lib.scheduler import rdb_requests 16from autotest_lib.scheduler import rdb_testing_utils 17from autotest_lib.scheduler import rdb_utils 18 19 20from django.core import exceptions as django_exceptions 21from django.db.models import fields 22 23 24class RDBBaseRequestHandlerTests(unittest.TestCase): 25 """Base Request Handler Unittests.""" 26 27 def setUp(self): 28 self.handler = rdb.BaseHostRequestHandler() 29 self.handler.host_query_manager = mock.MagicMock() 30 self.update_manager = rdb_requests.BaseHostRequestManager( 31 rdb_requests.UpdateHostRequest, rdb.update_hosts) 32 self.get_hosts_manager = rdb_requests.BaseHostRequestManager( 33 rdb_requests.HostRequest, rdb.get_hosts) 34 35 36 def tearDown(self): 37 self.handler.host_query_manager.reset_mock() 38 39 40 def testResponseMapUpdate(self): 41 """Test response map behaviour. 42 43 Test that either adding an empty response against a request, or 2 44 responses for the same request will raise an exception. 45 """ 46 self.get_hosts_manager.add_request(host_id=1) 47 request = self.get_hosts_manager.request_queue[0] 48 response = [] 49 self.assertRaises( 50 rdb_utils.RDBException, self.handler.update_response_map, 51 *(request, response)) 52 response.append(rdb_testing_utils.FakeHost(hostname='host', host_id=1)) 53 self.handler.update_response_map(request, response) 54 self.assertRaises( 55 rdb_utils.RDBException, self.handler.update_response_map, 56 *(request, response)) 57 58 59 def testResponseMapChecking(self): 60 """Test response map sanity check. 61 62 Test that adding the same RDBHostServerWrapper for 2 requests will 63 raise an exception. 64 """ 65 # Assign the same host to 2 requests and check for exceptions. 66 self.get_hosts_manager.add_request(host_id=1) 67 self.get_hosts_manager.add_request(host_id=2) 68 request_1 = self.get_hosts_manager.request_queue[0] 69 request_2 = self.get_hosts_manager.request_queue[1] 70 response = [rdb_testing_utils.FakeHost(hostname='host', host_id=1)] 71 72 self.handler.update_response_map(request_1, response) 73 self.handler.update_response_map(request_2, response) 74 self.assertRaises( 75 rdb_utils.RDBException, self.handler.get_response) 76 77 # Assign the same exception to 2 requests and make sure there isn't a 78 # an exception, then check that the response returned is the 79 # exception_string and not the exception itself. 80 self.handler.response_map = {} 81 exception_string = 'This is an exception' 82 response = [rdb_utils.RDBException(exception_string)] 83 self.handler.update_response_map(request_1, response) 84 self.handler.update_response_map(request_2, response) 85 for response in self.handler.get_response().values(): 86 self.assertTrue(response[0] == exception_string) 87 88 89 def testBatchGetHosts(self): 90 """Test getting hosts. 91 92 Verify that: 93 1. We actually call get_hosts on the query_manager for a 94 batched_get_hosts request. 95 2. The hosts returned are matched up correctly with requests, 96 and each request gets exactly one response. 97 3. The hosts returned have all the fields needed to create an 98 RDBClientHostWrapper, in spite of having gone through the 99 to_wire process of serialization in get_response. 100 """ 101 fake_hosts = [] 102 for host_id in range(1, 4): 103 self.get_hosts_manager.add_request(host_id=host_id) 104 fake_hosts.append( 105 rdb_testing_utils.FakeHost('host%s'%host_id, host_id)) 106 self.handler.host_query_manager.get_hosts = mock.MagicMock( 107 return_value=fake_hosts) 108 self.handler.batch_get_hosts(self.get_hosts_manager.request_queue) 109 for request, hosts in self.handler.get_response().iteritems(): 110 self.assertTrue(len(hosts) == 1) 111 client_host = rdb_hosts.RDBClientHostWrapper(**hosts[0]) 112 self.assertTrue(request.host_id == client_host.id) 113 114 115 def testSingleUpdateRequest(self): 116 """Test that a single host update request hits the query manager.""" 117 payload = {'status': 'Ready'} 118 host_id = 10 119 self.update_manager.add_request(host_id=host_id, payload=payload) 120 self.handler.update_hosts(self.update_manager.request_queue) 121 self.handler.host_query_manager.update_hosts.assert_called_once_with( 122 [host_id], **payload) 123 124 125 def testDedupingSameHostRequests(self): 126 """Test same host 2 updates deduping.""" 127 payload_1 = {'status': 'Ready'} 128 payload_2 = {'locked': True} 129 host_id = 10 130 self.update_manager.add_request(host_id=host_id, payload=payload_1) 131 self.update_manager.add_request(host_id=host_id, payload=payload_2) 132 self.handler.update_hosts(self.update_manager.request_queue) 133 self.handler.host_query_manager.update_hosts.assert_called_once_with( 134 [host_id], **dict(payload_1.items() + payload_2.items())) 135 136 137 def testLastUpdateWins(self): 138 """Test 2 updates to the same row x column.""" 139 payload_1 = {'status': 'foobar'} 140 payload_2 = {'status': 'Ready'} 141 host_id = 10 142 self.update_manager.add_request(host_id=host_id, payload=payload_1) 143 self.update_manager.add_request(host_id=host_id, payload=payload_2) 144 self.handler.update_hosts(self.update_manager.request_queue) 145 self.handler.host_query_manager.update_hosts.assert_called_once_with( 146 [host_id], **payload_2) 147 148 149 def testDedupingSamePayloadRequests(self): 150 """Test same payload for 2 hosts only hits the db once.""" 151 payload = {'status': 'Ready'} 152 host_1_id = 10 153 host_2_id = 20 154 self.update_manager.add_request(host_id=host_1_id, payload=payload) 155 self.update_manager.add_request(host_id=host_2_id, payload=payload) 156 self.handler.update_hosts(self.update_manager.request_queue) 157 self.handler.host_query_manager.update_hosts.assert_called_once_with( 158 [host_1_id, host_2_id], **payload) 159 160 161 def testUpdateException(self): 162 """Test update exception handling. 163 164 1. An exception raised while processing one update shouldn't prevent 165 the others. 166 2. The exception shold get serialized as a string and returned via the 167 response map. 168 """ 169 payload = {'status': 'Ready'} 170 exception_msg = 'Bad Field' 171 exception_types = [django_exceptions.FieldError, 172 fields.FieldDoesNotExist] 173 self.update_manager.add_request(host_id=11, payload=payload) 174 self.update_manager.add_request(host_id=10, payload=payload) 175 mock_query_manager = self.handler.host_query_manager 176 177 for e, request in zip( 178 exception_types, self.update_manager.request_queue): 179 mock_query_manager.update_hosts.side_effect = e(exception_msg) 180 self.handler.update_hosts([request]) 181 182 response = self.handler.get_response() 183 for request in self.update_manager.request_queue: 184 self.assertTrue(exception_msg in response.get(request)) 185 186 187class QueryManagerTests(unittest.TestCase, 188 frontend_test_utils.FrontendTestMixin): 189 """Query Manager Tests.""" 190 191 def setUp(self): 192 self.db_helper = rdb_testing_utils.DBHelper() 193 self._database = self.db_helper.database 194 195 # Runs syncdb setting up initial database conditions 196 self._frontend_common_setup() 197 self.available_hosts_query_manager = rdb.AvailableHostQueryManager() 198 self.all_hosts_query_manager = rdb.BaseHostQueryManager() 199 200 201 def tearDown(self): 202 self._database.disconnect() 203 self._frontend_common_teardown() 204 205 206 def testFindHosts(self): 207 """Test finding hosts. 208 209 Tests that we can only find unleased hosts through the 210 available_hosts_query_manager. 211 """ 212 deps = set(['a', 'b']) 213 acls = set(['a']) 214 db_host = self.db_helper.create_host( 215 name='h1', deps=deps, acls=acls, leased=1) 216 hosts = self.all_hosts_query_manager.find_hosts( 217 deps=[lable.id for lable in db_host.labels.all()], 218 acls=[aclgroup.id for aclgroup in db_host.aclgroup_set.all()]) 219 self.assertTrue(type(hosts) == list and len(hosts) == 1) 220 hosts = self.available_hosts_query_manager.find_hosts( 221 deps=[lable.id for lable in db_host.labels.all()], 222 acls=[aclgroup.id for aclgroup in db_host.aclgroup_set.all()]) 223 # We should get an empty list if there are no matching hosts, not a 224 # QuerySet or None. 225 self.assertTrue(len(hosts) == 0) 226 227 228 def testUpdateHosts(self): 229 """Test updating hosts. 230 231 Test that we can only update unleased hosts through the 232 available_hosts_query_manager. 233 """ 234 deps = set(['a', 'b']) 235 acls = set(['a']) 236 db_host = self.db_helper.create_host( 237 name='h1', deps=deps, acls=acls, leased=1) 238 # Confirm that the available_hosts_manager can't see the leased host. 239 self.assertTrue( 240 len(self.available_hosts_query_manager.get_hosts( 241 [db_host.id])) == 0) 242 243 # Confirm that the available_hosts_manager can't update a leased host. 244 # Also confirm that the general query manager Can see the leased host. 245 self.available_hosts_query_manager.update_hosts( 246 [db_host.id], **{'leased': 0}) 247 hosts = self.all_hosts_query_manager.get_hosts([db_host.id]) 248 self.assertTrue(len(hosts) == 1 and hosts[0].leased) 249 250 # Confirm that we can infact update the leased bit on the host. 251 self.all_hosts_query_manager.update_hosts( 252 [hosts[0].id], **{'leased': 0}) 253 hosts = self.all_hosts_query_manager.get_hosts([hosts[0].id]) 254 self.assertTrue(len(hosts) == 1 and not hosts[0].leased) 255 256 257if __name__ == '__main__': 258 unittest.main() 259