def test_user_or_global_default_drone_set_no_user(self):
expected = self._setup_test_user_or_global_default_drone_set()
actual = agent_task.AgentTask()._user_or_global_default_drone_set(
None, None)
self.assertEqual(expected, actual)
self.god.check_playback()
def _setup_drones(self):
self.god.stub_function(models.DroneSet, 'drone_sets_enabled')
models.DroneSet.drone_sets_enabled.expect_call().and_return(True)
drones = []
for x in xrange(4):
drones.append(models.Drone.objects.create(hostname=str(x)))
drone_set_1 = models.DroneSet.objects.create(name='1')
drone_set_1.drones.add(*drones[0:2])
drone_set_2 = models.DroneSet.objects.create(name='2')
drone_set_2.drones.add(*drones[2:4])
drone_set_3 = models.DroneSet.objects.create(name='3')
job_1 = self._create_job_simple([self.hosts[0].id],
drone_set=drone_set_1)
job_2 = self._create_job_simple([self.hosts[0].id],
drone_set=drone_set_2)
job_3 = self._create_job_simple([self.hosts[0].id],
drone_set=drone_set_3)
job_4 = self._create_job_simple([self.hosts[0].id])
job_4.drone_set = None
job_4.save()
hqe_1 = job_1.hostqueueentry_set.all()[0]
hqe_2 = job_2.hostqueueentry_set.all()[0]
hqe_3 = job_3.hostqueueentry_set.all()[0]
hqe_4 = job_4.hostqueueentry_set.all()[0]
return (hqe_1, hqe_2, hqe_3, hqe_4), agent_task.AgentTask()
def test_user_or_global_default_drone_set_no_user_drone_set(self):
class MockUser(object):
drone_set = None
login = None
expected = self._setup_test_user_or_global_default_drone_set()
actual = agent_task.AgentTask()._user_or_global_default_drone_set(
None, MockUser())
self.assertEqual(expected, actual)
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()
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_user_or_global_default_drone_set(self):
expected = object()
class MockDroneSet(object):
def get_drone_hostnames(self):
return expected
class MockUser(object):
drone_set = MockDroneSet()
self._setup_test_user_or_global_default_drone_set()
actual = agent_task.AgentTask()._user_or_global_default_drone_set(
None, MockUser())
self.assertEqual(expected, actual)
self.god.check_playback()