1 #!/usr/bin/python 2 3 import common 4 from autotest_lib.frontend import setup_django_environment 5 from autotest_lib.frontend.afe import frontend_test_utils 6 from autotest_lib.client.common_lib.test_utils import unittest 7 from autotest_lib.frontend.afe import models 8 from autotest_lib.scheduler import agent_task 9 from autotest_lib.server import system_utils 10 11 12 class RestrictedSubnetTest(unittest.TestCase, 13 frontend_test_utils.FrontendTestMixin): 14 """Test server election based on restricted subnet setting. 15 """ 16 17 DRONE_IN_RESTRICTED_SUBNET = '192.168.0.9' 18 DRONE_NOT_IN_RESTRICTED_SUBNET = '127.0.0.9' 19 HOST_IN_RESTRICTED_SUBNET = '192.168.0.3' 20 HOST_NOT_IN_RESTRICTED_SUBNET = '127.0.0.3' 21 RESTRICTED_SUBNETS = [('192.168.0.1', 16)] 22 23 def setUp(self): 24 self._drones = [self.DRONE_IN_RESTRICTED_SUBNET, 25 self.DRONE_NOT_IN_RESTRICTED_SUBNET] 26 system_utils.DroneCache.unrestricted_drones = None 27 system_utils.DroneCache.drone_ip_map = None 28 self._frontend_common_setup() 29 30 31 def tearDown(self): 32 self._frontend_common_teardown() 33 34 35 def test_get_drone_hostnames_allowed_with_restricted_subnet(self): 36 """Test method get_drone_hostnames_allowed work as expected when 37 restricted subnet is set, and host is inside restricted subnet. 38 """ 39 self.god.stub_function(system_utils, 'get_drones') 40 system_utils.get_drones.expect_call().and_return(self._drones) 41 self.god.stub_function(models.DroneSet, 'drone_sets_enabled') 42 models.DroneSet.drone_sets_enabled.expect_call().and_return(False) 43 44 task = agent_task.AgentTask() 45 task.hostnames = {1: self.HOST_IN_RESTRICTED_SUBNET} 46 self.assertEqual( 47 set([self.DRONE_IN_RESTRICTED_SUBNET]), 48 task.get_drone_hostnames_allowed(self.RESTRICTED_SUBNETS, True)) 49 self.god.check_playback() 50 51 52 def test_get_drone_hostnames_allowed_not_in_restricted_subnet(self): 53 """Test method get_drone_hostnames_allowed work as expected when 54 restricted subnet is set, and host is not in restricted subnet. 55 """ 56 self.god.stub_function(system_utils, 'get_drones') 57 system_utils.get_drones.expect_call().and_return(self._drones) 58 self.god.stub_function(models.DroneSet, 'drone_sets_enabled') 59 models.DroneSet.drone_sets_enabled.expect_call().and_return(False) 60 61 task = agent_task.AgentTask() 62 task.hostnames = {1: self.HOST_NOT_IN_RESTRICTED_SUBNET} 63 self.assertEqual( 64 set([self.DRONE_NOT_IN_RESTRICTED_SUBNET]), 65 task.get_drone_hostnames_allowed(self.RESTRICTED_SUBNETS, True)) 66 self.god.check_playback() 67 68 69 def test_get_drone_hostnames_allowed_in_mixed_subnet(self): 70 """Test method get_drone_hostnames_allowed work as expected when 71 restricted subnet is set, and hosts are distributed across restricted 72 subnet and unrestricted subnet. 73 """ 74 task = agent_task.AgentTask() 75 task.hostnames = {1: self.HOST_NOT_IN_RESTRICTED_SUBNET, 76 2: self.HOST_IN_RESTRICTED_SUBNET} 77 self.assertEqual( 78 set(), 79 task.get_drone_hostnames_allowed(self.RESTRICTED_SUBNETS, True)) 80 self.god.check_playback() 81 82 83 if __name__ == '__main__': 84 unittest.main() 85
