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