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