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