init repo

ronnie

2022-10-14 1504bb53e29d3d46222c0b3ea994fc494b48e153

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