1#!/usr/bin/python
2
3import unittest
4
5import common
6from autotest_lib.frontend import setup_django_environment
7from autotest_lib.frontend.afe import frontend_test_utils
8from autotest_lib.frontend.afe import models
9from autotest_lib.scheduler import agent_task
10from autotest_lib.server import system_utils
11
12
13class 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
84if __name__ == '__main__':
85    unittest.main()
86