1 import datetime 2 3 import common 4 from autotest_lib.frontend import setup_test_environment 5 from autotest_lib.frontend import thread_local 6 from autotest_lib.frontend.afe import models, model_attributes 7 from autotest_lib.client.common_lib import global_config 8 from autotest_lib.client.common_lib.test_utils import mock 9 10 class FrontendTestMixin(object): 11 # pylint: disable=missing-docstring 12 def _fill_in_test_data(self): 13 """Populate the test database with some hosts and labels.""" 14 if models.DroneSet.drone_sets_enabled(): 15 models.DroneSet.objects.create( 16 name=models.DroneSet.default_drone_set_name()) 17 18 acl_group = models.AclGroup.objects.create(name='my_acl') 19 acl_group.users.add(models.User.current_user()) 20 21 self.hosts = [models.Host.objects.create(hostname=hostname) 22 for hostname in 23 ('host1', 'host2', 'host3', 'host4', 'host5', 'host6', 24 'host7', 'host8', 'host9')] 25 26 acl_group.hosts = self.hosts 27 models.AclGroup.smart_get('Everyone').hosts = [] 28 29 self.labels = [models.Label.objects.create(name=name) for name in 30 ('label1', 'label2', 'label3', 31 'label6', 'label7', 'unused')] 32 33 platform = models.Label.objects.create(name='myplatform', platform=True) 34 for host in self.hosts: 35 host.labels.add(platform) 36 37 self.label1, self.label2, self.label3, self.label6, self.label7, _ \ 38 = self.labels 39 40 self.label3.only_if_needed = True 41 self.label3.save() 42 self.hosts[0].labels.add(self.label1) 43 self.hosts[1].labels.add(self.label2) 44 for hostnum in xrange(4,7): # host5..host7 45 self.hosts[hostnum].labels.add(self.label6) 46 self.hosts[6].labels.add(self.label7) 47 for hostnum in xrange(7,9): # host8..host9 48 self.hosts[hostnum].labels.add(self.label6) 49 self.hosts[hostnum].labels.add(self.label7) 50 51 52 def _frontend_common_setup(self, fill_data=True, setup_tables=True): 53 self.god = mock.mock_god(ut=self) 54 if setup_tables: 55 setup_test_environment.set_up() 56 global_config.global_config.override_config_value( 57 'SERVER', 'rpc_logging', 'False') 58 if fill_data and setup_tables: 59 self._fill_in_test_data() 60 61 62 def _frontend_common_teardown(self): 63 setup_test_environment.tear_down() 64 thread_local.set_user(None) 65 self.god.unstub_all() 66 67 68 def _create_job(self, hosts=[], metahosts=[], priority=0, active=False, 69 synchronous=False, hostless=False, 70 drone_set=None, control_file='control', 71 owner='autotest_system', parent_job_id=None, 72 shard=None): 73 """ 74 Create a job row in the test database. 75 76 @param hosts - A list of explicit host ids for this job to be 77 scheduled on. 78 @param metahosts - A list of label ids for each host that this job 79 should be scheduled on (meta host scheduling). 80 @param priority - The job priority (integer). 81 @param active - bool, mark this job as running or not in the database? 82 @param synchronous - bool, if True use synch_count=2 otherwise use 83 synch_count=1. 84 @param hostless - if True, this job is intended to be hostless (in that 85 case, hosts, and metahosts must all be empty) 86 @param owner - The owner of the job. Aclgroups from which a job can 87 acquire hosts change with the aclgroups of the owners. 88 @param parent_job_id - The id of a parent_job. If a job with the id 89 doesn't already exist one will be created. 90 @param shard - shard object to assign the job to. 91 92 @raises model.DoesNotExist: If parent_job_id is specified but a job with 93 id=parent_job_id does not exist. 94 95 @returns A Django frontend.afe.models.Job instance. 96 """ 97 if not drone_set: 98 drone_set = (models.DroneSet.default_drone_set_name() 99 and models.DroneSet.get_default()) 100 101 synch_count = synchronous and 2 or 1 102 created_on = datetime.datetime(2008, 1, 1) 103 status = models.HostQueueEntry.Status.QUEUED 104 if active: 105 status = models.HostQueueEntry.Status.RUNNING 106 107 parent_job = (models.Job.objects.get(id=parent_job_id) 108 if parent_job_id else None) 109 job = models.Job.objects.create( 110 name='test', owner=owner, priority=priority, 111 synch_count=synch_count, created_on=created_on, 112 reboot_before=model_attributes.RebootBefore.NEVER, 113 drone_set=drone_set, control_file=control_file, 114 parent_job=parent_job, require_ssp=None, 115 shard=shard) 116 117 # Update the job's dependencies to include the metahost. 118 for metahost_label in metahosts: 119 dep = models.Label.objects.get(id=metahost_label) 120 job.dependency_labels.add(dep) 121 122 for host_id in hosts: 123 models.HostQueueEntry.objects.create(job=job, host_id=host_id, 124 status=status) 125 models.IneligibleHostQueue.objects.create(job=job, host_id=host_id) 126 for label_id in metahosts: 127 models.HostQueueEntry.objects.create(job=job, meta_host_id=label_id, 128 status=status) 129 130 if hostless: 131 assert not (hosts or metahosts) 132 models.HostQueueEntry.objects.create(job=job, status=status) 133 return job 134 135 136 def _create_job_simple(self, hosts, use_metahost=False, 137 priority=0, active=False, drone_set=None, 138 parent_job_id=None): 139 """An alternative interface to _create_job""" 140 args = {'hosts' : [], 'metahosts' : []} 141 if use_metahost: 142 args['metahosts'] = hosts 143 else: 144 args['hosts'] = hosts 145 return self._create_job( 146 priority=priority, active=active, drone_set=drone_set, 147 parent_job_id=parent_job_id, **args) 148