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