Don Garrett | 580717f | 2015-07-24 14:11:22 -0700 | [diff] [blame] | 1 | #!/usr/bin/python |
| 2 | # |
Prashanth B | b474fdf | 2014-04-03 16:05:38 -0700 | [diff] [blame] | 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 | |
| 7 | import unittest |
| 8 | |
| 9 | import common |
Prashanth B | b474fdf | 2014-04-03 16:05:38 -0700 | [diff] [blame] | 10 | from autotest_lib.frontend import setup_django_environment |
| 11 | from autotest_lib.frontend.afe import frontend_test_utils |
| 12 | from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models |
| 13 | from autotest_lib.scheduler import rdb_hosts |
Prashanth B | 2d8047e | 2014-04-27 18:54:47 -0700 | [diff] [blame] | 14 | from autotest_lib.scheduler import rdb_testing_utils |
Prashanth B | b474fdf | 2014-04-03 16:05:38 -0700 | [diff] [blame] | 15 | from autotest_lib.scheduler import rdb_utils |
| 16 | |
| 17 | |
| 18 | class RDBHostTests(unittest.TestCase, frontend_test_utils.FrontendTestMixin): |
| 19 | """Unittests for RDBHost objects.""" |
| 20 | |
| 21 | def setUp(self): |
Prashanth B | 2d8047e | 2014-04-27 18:54:47 -0700 | [diff] [blame] | 22 | self.db_helper = rdb_testing_utils.DBHelper() |
Prashanth B | b474fdf | 2014-04-03 16:05:38 -0700 | [diff] [blame] | 23 | self._database = self.db_helper.database |
| 24 | # Runs syncdb setting up initial database conditions |
| 25 | self._frontend_common_setup() |
| 26 | |
| 27 | |
| 28 | def tearDown(self): |
| 29 | self._database.disconnect() |
| 30 | self._frontend_common_teardown() |
| 31 | |
| 32 | |
| 33 | def testWireFormat(self): |
| 34 | """Test that we can create a client host with the server host's fields. |
| 35 | |
| 36 | Get the wire_format fields of an RDBServerHostWrapper and use them to |
| 37 | create an RDBClientHostWrapper. |
| 38 | |
| 39 | @raises AssertionError: If the labels and acls don't match up after |
| 40 | going through the complete wire_format conversion, of the bare |
| 41 | wire_format conversion also converts labels and acls. |
| 42 | @raises RDBException: If some critical fields were lost during |
| 43 | wire_format conversion, as we won't be able to construct the |
| 44 | RDBClientHostWrapper. |
| 45 | """ |
| 46 | labels = set(['a', 'b', 'c']) |
| 47 | acls = set(['d', 'e']) |
| 48 | server_host = rdb_hosts.RDBServerHostWrapper( |
| 49 | self.db_helper.create_host('h1', deps=labels, acls=acls)) |
| 50 | acl_ids = set([aclgroup.id for aclgroup in |
| 51 | self.db_helper.get_acls(name__in=acls)]) |
| 52 | label_ids = set([label.id for label in |
| 53 | self.db_helper.get_labels(name__in=labels)]) |
| 54 | |
| 55 | # The RDBServerHostWrapper keeps ids of labels/acls to perform |
| 56 | # comparison operations within the rdb, but converts labels to |
| 57 | # strings because this is the format the scheduler expects them in. |
| 58 | self.assertTrue(set(server_host.labels) == label_ids and |
| 59 | set(server_host.acls) == acl_ids) |
| 60 | |
| 61 | formatted_server_host = server_host.wire_format() |
| 62 | client_host = rdb_hosts.RDBClientHostWrapper(**formatted_server_host) |
| 63 | self.assertTrue(set(client_host.labels) == labels and |
| 64 | set(client_host.acls) == acl_ids) |
| 65 | bare_formatted_server_host = server_host.wire_format( |
| 66 | unwrap_foreign_keys=False) |
| 67 | self.assertTrue(bare_formatted_server_host.get('labels') is None and |
| 68 | bare_formatted_server_host.get('acls') is None) |
| 69 | |
| 70 | |
| 71 | def testLeasing(self): |
| 72 | """Test that leasing a leased host raises an exception. |
| 73 | |
| 74 | @raises AssertionError: If double leasing a host doesn't raise |
| 75 | an RDBException, or the leased bits are not set after the |
| 76 | first attempt at leasing it. |
| 77 | @raises RDBException: If the host is created with the leased bit set. |
| 78 | """ |
| 79 | hostname = 'h1' |
| 80 | server_host = rdb_hosts.RDBServerHostWrapper( |
| 81 | self.db_helper.create_host(hostname)) |
| 82 | server_host.lease() |
| 83 | host = self.db_helper.get_host(hostname=hostname)[0] |
| 84 | self.assertTrue(host.leased and server_host.leased) |
| 85 | self.assertRaises(rdb_utils.RDBException, server_host.lease) |
| 86 | |
| 87 | |
| 88 | def testPlatformAndLabels(self): |
| 89 | """Test that a client host returns the right platform and labels. |
| 90 | |
| 91 | @raises AssertionError: If client host cannot return the right platform |
| 92 | and labels. |
| 93 | """ |
| 94 | platform_name = 'x86' |
| 95 | label_names = ['a', 'b'] |
| 96 | self.db_helper.create_label(name=platform_name, platform=True) |
| 97 | server_host = rdb_hosts.RDBServerHostWrapper( |
| 98 | self.db_helper.create_host( |
| 99 | 'h1', deps=set(label_names + [platform_name]))) |
| 100 | client_host = rdb_hosts.RDBClientHostWrapper( |
| 101 | **server_host.wire_format()) |
| 102 | platform, labels = client_host.platform_and_labels() |
| 103 | self.assertTrue(platform == platform_name) |
| 104 | self.assertTrue(set(labels) == set(label_names)) |
| 105 | |
| 106 | |
| 107 | def testClientUpdateSave(self): |
| 108 | """Test that a client host is capable of saving its attributes. |
| 109 | |
| 110 | Create a client host, set its attributes and verify that the attributes |
| 111 | are saved properly by recreating a server host and checking them. |
| 112 | |
| 113 | @raises AssertionError: If the server host has the wrong attributes. |
| 114 | """ |
| 115 | hostname = 'h1' |
| 116 | db_host = self.db_helper.create_host(hostname, leased=True) |
| 117 | server_host_dict = rdb_hosts.RDBServerHostWrapper(db_host).wire_format() |
| 118 | client_host = rdb_hosts.RDBClientHostWrapper(**server_host_dict) |
| 119 | |
| 120 | host_data = {'hostname': hostname, 'id': db_host.id} |
| 121 | default_values = rdb_models.AbstractHostModel.provide_default_values( |
| 122 | host_data) |
| 123 | for k, v in default_values.iteritems(): |
| 124 | self.assertTrue(server_host_dict[k] == v) |
| 125 | |
| 126 | updated_client_fields = { |
| 127 | 'locked': True, |
| 128 | 'leased': False, |
| 129 | 'status': 'FakeStatus', |
| 130 | 'invalid': True, |
| 131 | 'protection': 1, |
| 132 | 'dirty': True, |
| 133 | } |
| 134 | client_host.__dict__.update(updated_client_fields) |
| 135 | client_host.save() |
| 136 | |
| 137 | updated_server_host = rdb_hosts.RDBServerHostWrapper( |
| 138 | self.db_helper.get_host(hostname=hostname)[0]).wire_format() |
| 139 | for k, v in updated_client_fields.iteritems(): |
| 140 | self.assertTrue(updated_server_host[k] == v) |
| 141 | |
| 142 | |
| 143 | def testUpdateField(self): |
| 144 | """Test that update field on the client host works as expected. |
| 145 | |
| 146 | @raises AssertionError: If a bad update is processed without an |
| 147 | exception, of a good update isn't processed as expected. |
| 148 | """ |
| 149 | hostname = 'h1' |
| 150 | db_host = self.db_helper.create_host(hostname, dirty=False) |
| 151 | server_host_dict = rdb_hosts.RDBServerHostWrapper(db_host).wire_format() |
| 152 | client_host = rdb_hosts.RDBClientHostWrapper(**server_host_dict) |
| 153 | self.assertRaises(rdb_utils.RDBException, client_host.update_field, |
| 154 | *('id', 'fakeid')) |
| 155 | self.assertRaises(rdb_utils.RDBException, client_host.update_field, |
| 156 | *('Nonexist', 'Nonexist')) |
| 157 | client_host.update_field('dirty', True) |
| 158 | self.assertTrue( |
| 159 | self.db_helper.get_host(hostname=hostname)[0].dirty == True and |
| 160 | client_host.dirty == True) |
| 161 | new_status = 'newstatus' |
| 162 | client_host.set_status(new_status) |
| 163 | self.assertTrue( |
| 164 | self.db_helper.get_host(hostname=hostname)[0].status == |
| 165 | new_status and client_host.status == new_status) |
| 166 | |
Don Garrett | 580717f | 2015-07-24 14:11:22 -0700 | [diff] [blame] | 167 | |
| 168 | if __name__ == '__main__': |
| 169 | unittest.main() |