| #!/usr/bin/python |
| # |
| # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| import mock |
| import unittest |
| |
| import common |
| from autotest_lib.frontend import setup_django_lite_environment |
| from autotest_lib.frontend.afe import frontend_test_utils |
| from autotest_lib.scheduler import rdb |
| from autotest_lib.scheduler import rdb_hosts |
| from autotest_lib.scheduler import rdb_requests |
| from autotest_lib.scheduler import rdb_testing_utils |
| from autotest_lib.scheduler import rdb_utils |
| |
| |
| from django.core import exceptions as django_exceptions |
| from django.db.models import fields |
| |
| |
| class RDBBaseRequestHandlerTests(unittest.TestCase): |
| """Base Request Handler Unittests.""" |
| |
| def setUp(self): |
| self.handler = rdb.BaseHostRequestHandler() |
| self.handler.host_query_manager = mock.MagicMock() |
| self.update_manager = rdb_requests.BaseHostRequestManager( |
| rdb_requests.UpdateHostRequest, rdb.update_hosts) |
| self.get_hosts_manager = rdb_requests.BaseHostRequestManager( |
| rdb_requests.HostRequest, rdb.get_hosts) |
| |
| |
| def tearDown(self): |
| self.handler.host_query_manager.reset_mock() |
| |
| |
| def testResponseMapUpdate(self): |
| """Test response map behaviour. |
| |
| Test that either adding an empty response against a request, or 2 |
| responses for the same request will raise an exception. |
| """ |
| self.get_hosts_manager.add_request(host_id=1) |
| request = self.get_hosts_manager.request_queue[0] |
| response = [] |
| self.assertRaises( |
| rdb_utils.RDBException, self.handler.update_response_map, |
| *(request, response)) |
| response.append(rdb_testing_utils.FakeHost(hostname='host', host_id=1)) |
| self.handler.update_response_map(request, response) |
| self.assertRaises( |
| rdb_utils.RDBException, self.handler.update_response_map, |
| *(request, response)) |
| |
| |
| def testResponseMapChecking(self): |
| """Test response map sanity check. |
| |
| Test that adding the same RDBHostServerWrapper for 2 requests will |
| raise an exception. |
| """ |
| # Assign the same host to 2 requests and check for exceptions. |
| self.get_hosts_manager.add_request(host_id=1) |
| self.get_hosts_manager.add_request(host_id=2) |
| request_1 = self.get_hosts_manager.request_queue[0] |
| request_2 = self.get_hosts_manager.request_queue[1] |
| response = [rdb_testing_utils.FakeHost(hostname='host', host_id=1)] |
| |
| self.handler.update_response_map(request_1, response) |
| self.handler.update_response_map(request_2, response) |
| self.assertRaises( |
| rdb_utils.RDBException, self.handler.get_response) |
| |
| # Assign the same exception to 2 requests and make sure there isn't a |
| # an exception, then check that the response returned is the |
| # exception_string and not the exception itself. |
| self.handler.response_map = {} |
| exception_string = 'This is an exception' |
| response = [rdb_utils.RDBException(exception_string)] |
| self.handler.update_response_map(request_1, response) |
| self.handler.update_response_map(request_2, response) |
| for response in self.handler.get_response().values(): |
| self.assertTrue(response[0] == exception_string) |
| |
| |
| def testBatchGetHosts(self): |
| """Test getting hosts. |
| |
| Verify that: |
| 1. We actually call get_hosts on the query_manager for a |
| batched_get_hosts request. |
| 2. The hosts returned are matched up correctly with requests, |
| and each request gets exactly one response. |
| 3. The hosts returned have all the fields needed to create an |
| RDBClientHostWrapper, in spite of having gone through the |
| to_wire process of serialization in get_response. |
| """ |
| fake_hosts = [] |
| for host_id in range(1, 4): |
| self.get_hosts_manager.add_request(host_id=host_id) |
| fake_hosts.append( |
| rdb_testing_utils.FakeHost('host%s'%host_id, host_id)) |
| self.handler.host_query_manager.get_hosts = mock.MagicMock( |
| return_value=fake_hosts) |
| self.handler.batch_get_hosts(self.get_hosts_manager.request_queue) |
| for request, hosts in self.handler.get_response().iteritems(): |
| self.assertTrue(len(hosts) == 1) |
| client_host = rdb_hosts.RDBClientHostWrapper(**hosts[0]) |
| self.assertTrue(request.host_id == client_host.id) |
| |
| |
| def testSingleUpdateRequest(self): |
| """Test that a single host update request hits the query manager.""" |
| payload = {'status': 'Ready'} |
| host_id = 10 |
| self.update_manager.add_request(host_id=host_id, payload=payload) |
| self.handler.update_hosts(self.update_manager.request_queue) |
| self.handler.host_query_manager.update_hosts.assert_called_once_with( |
| [host_id], **payload) |
| |
| |
| def testDedupingSameHostRequests(self): |
| """Test same host 2 updates deduping.""" |
| payload_1 = {'status': 'Ready'} |
| payload_2 = {'locked': True} |
| host_id = 10 |
| self.update_manager.add_request(host_id=host_id, payload=payload_1) |
| self.update_manager.add_request(host_id=host_id, payload=payload_2) |
| self.handler.update_hosts(self.update_manager.request_queue) |
| self.handler.host_query_manager.update_hosts.assert_called_once_with( |
| [host_id], **dict(payload_1.items() + payload_2.items())) |
| |
| |
| def testLastUpdateWins(self): |
| """Test 2 updates to the same row x column.""" |
| payload_1 = {'status': 'foobar'} |
| payload_2 = {'status': 'Ready'} |
| host_id = 10 |
| self.update_manager.add_request(host_id=host_id, payload=payload_1) |
| self.update_manager.add_request(host_id=host_id, payload=payload_2) |
| self.handler.update_hosts(self.update_manager.request_queue) |
| self.handler.host_query_manager.update_hosts.assert_called_once_with( |
| [host_id], **payload_2) |
| |
| |
| def testDedupingSamePayloadRequests(self): |
| """Test same payload for 2 hosts only hits the db once.""" |
| payload = {'status': 'Ready'} |
| host_1_id = 10 |
| host_2_id = 20 |
| self.update_manager.add_request(host_id=host_1_id, payload=payload) |
| self.update_manager.add_request(host_id=host_2_id, payload=payload) |
| self.handler.update_hosts(self.update_manager.request_queue) |
| self.handler.host_query_manager.update_hosts.assert_called_once_with( |
| [host_1_id, host_2_id], **payload) |
| |
| |
| def testUpdateException(self): |
| """Test update exception handling. |
| |
| 1. An exception raised while processing one update shouldn't prevent |
| the others. |
| 2. The exception shold get serialized as a string and returned via the |
| response map. |
| """ |
| payload = {'status': 'Ready'} |
| exception_msg = 'Bad Field' |
| exception_types = [django_exceptions.FieldError, |
| fields.FieldDoesNotExist] |
| self.update_manager.add_request(host_id=11, payload=payload) |
| self.update_manager.add_request(host_id=10, payload=payload) |
| mock_query_manager = self.handler.host_query_manager |
| |
| for e, request in zip( |
| exception_types, self.update_manager.request_queue): |
| mock_query_manager.update_hosts.side_effect = e(exception_msg) |
| self.handler.update_hosts([request]) |
| |
| response = self.handler.get_response() |
| for request in self.update_manager.request_queue: |
| self.assertTrue(exception_msg in response.get(request)) |
| |
| |
| class QueryManagerTests(unittest.TestCase, |
| frontend_test_utils.FrontendTestMixin): |
| """Query Manager Tests.""" |
| |
| def setUp(self): |
| self.db_helper = rdb_testing_utils.DBHelper() |
| self._database = self.db_helper.database |
| |
| # Runs syncdb setting up initial database conditions |
| self._frontend_common_setup() |
| self.available_hosts_query_manager = rdb.AvailableHostQueryManager() |
| self.all_hosts_query_manager = rdb.BaseHostQueryManager() |
| |
| |
| def tearDown(self): |
| self._database.disconnect() |
| self._frontend_common_teardown() |
| |
| |
| def testFindHosts(self): |
| """Test finding hosts. |
| |
| Tests that we can only find unleased hosts through the |
| available_hosts_query_manager. |
| """ |
| deps = set(['a', 'b']) |
| acls = set(['a']) |
| db_host = self.db_helper.create_host( |
| name='h1', deps=deps, acls=acls, leased=1) |
| hosts = self.all_hosts_query_manager.find_hosts( |
| deps=[lable.id for lable in db_host.labels.all()], |
| acls=[aclgroup.id for aclgroup in db_host.aclgroup_set.all()]) |
| self.assertTrue(type(hosts) == list and len(hosts) == 1) |
| hosts = self.available_hosts_query_manager.find_hosts( |
| deps=[lable.id for lable in db_host.labels.all()], |
| acls=[aclgroup.id for aclgroup in db_host.aclgroup_set.all()]) |
| # We should get an empty list if there are no matching hosts, not a |
| # QuerySet or None. |
| self.assertTrue(len(hosts) == 0) |
| |
| |
| def testUpdateHosts(self): |
| """Test updating hosts. |
| |
| Test that we can only update unleased hosts through the |
| available_hosts_query_manager. |
| """ |
| deps = set(['a', 'b']) |
| acls = set(['a']) |
| db_host = self.db_helper.create_host( |
| name='h1', deps=deps, acls=acls, leased=1) |
| # Confirm that the available_hosts_manager can't see the leased host. |
| self.assertTrue( |
| len(self.available_hosts_query_manager.get_hosts( |
| [db_host.id])) == 0) |
| |
| # Confirm that the available_hosts_manager can't update a leased host. |
| # Also confirm that the general query manager Can see the leased host. |
| self.available_hosts_query_manager.update_hosts( |
| [db_host.id], **{'leased': 0}) |
| hosts = self.all_hosts_query_manager.get_hosts([db_host.id]) |
| self.assertTrue(len(hosts) == 1 and hosts[0].leased) |
| |
| # Confirm that we can infact update the leased bit on the host. |
| self.all_hosts_query_manager.update_hosts( |
| [hosts[0].id], **{'leased': 0}) |
| hosts = self.all_hosts_query_manager.get_hosts([hosts[0].id]) |
| self.assertTrue(len(hosts) == 1 and not hosts[0].leased) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |