1import datetime 2 3import common 4from autotest_lib.frontend import setup_test_environment 5from autotest_lib.frontend import thread_local 6from autotest_lib.frontend.afe import models, model_attributes 7from autotest_lib.client.common_lib import global_config 8from autotest_lib.client.common_lib.test_utils import mock 9 10class 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.labels.append(models.Label.objects.create(name='static')) 41 self.replaced_labels = [models.ReplacedLabel.objects.create( 42 label_id=self.labels[-1].id)] 43 44 self.label3.only_if_needed = True 45 self.label3.save() 46 self.hosts[0].labels.add(self.label1) 47 self.hosts[1].labels.add(self.label2) 48 for hostnum in xrange(4,7): # host5..host7 49 self.hosts[hostnum].labels.add(self.label6) 50 self.hosts[6].labels.add(self.label7) 51 for hostnum in xrange(7,9): # host8..host9 52 self.hosts[hostnum].labels.add(self.label6) 53 self.hosts[hostnum].labels.add(self.label7) 54 55 56 def _frontend_common_setup(self, fill_data=True, setup_tables=True): 57 self.god = mock.mock_god(ut=self) 58 if setup_tables: 59 setup_test_environment.set_up() 60 global_config.global_config.override_config_value( 61 'SERVER', 'rpc_logging', 'False') 62 if fill_data and setup_tables: 63 self._fill_in_test_data() 64 65 66 def _frontend_common_teardown(self): 67 setup_test_environment.tear_down() 68 thread_local.set_user(None) 69 self.god.unstub_all() 70 71 72 def _set_static_attribute(self, host, attribute, value): 73 """Set static attribute for a host. 74 75 It ensures that all static attributes have a corresponding 76 entry in afe_host_attributes. 77 """ 78 # Get or create the reference object in afe_host_attributes. 79 model, args = host._get_attribute_model_and_args(attribute) 80 model.objects.get_or_create(**args) 81 82 attribute_model, get_args = host._get_static_attribute_model_and_args( 83 attribute) 84 attribute_object, _ = attribute_model.objects.get_or_create(**get_args) 85 attribute_object.value = value 86 attribute_object.save() 87 88 89 def _create_job(self, hosts=[], metahosts=[], priority=0, active=False, 90 synchronous=False, hostless=False, 91 drone_set=None, control_file='control', 92 owner='autotest_system', parent_job_id=None, 93 shard=None): 94 """ 95 Create a job row in the test database. 96 97 @param hosts - A list of explicit host ids for this job to be 98 scheduled on. 99 @param metahosts - A list of label ids for each host that this job 100 should be scheduled on (meta host scheduling). 101 @param priority - The job priority (integer). 102 @param active - bool, mark this job as running or not in the database? 103 @param synchronous - bool, if True use synch_count=2 otherwise use 104 synch_count=1. 105 @param hostless - if True, this job is intended to be hostless (in that 106 case, hosts, and metahosts must all be empty) 107 @param owner - The owner of the job. Aclgroups from which a job can 108 acquire hosts change with the aclgroups of the owners. 109 @param parent_job_id - The id of a parent_job. If a job with the id 110 doesn't already exist one will be created. 111 @param shard - shard object to assign the job to. 112 113 @raises model.DoesNotExist: If parent_job_id is specified but a job with 114 id=parent_job_id does not exist. 115 116 @returns A Django frontend.afe.models.Job instance. 117 """ 118 if not drone_set: 119 drone_set = (models.DroneSet.default_drone_set_name() 120 and models.DroneSet.get_default()) 121 122 synch_count = synchronous and 2 or 1 123 created_on = datetime.datetime(2008, 1, 1) 124 status = models.HostQueueEntry.Status.QUEUED 125 if active: 126 status = models.HostQueueEntry.Status.RUNNING 127 128 parent_job = (models.Job.objects.get(id=parent_job_id) 129 if parent_job_id else None) 130 job = models.Job.objects.create( 131 name='test', owner=owner, priority=priority, 132 synch_count=synch_count, created_on=created_on, 133 reboot_before=model_attributes.RebootBefore.NEVER, 134 drone_set=drone_set, control_file=control_file, 135 parent_job=parent_job, require_ssp=None, 136 shard=shard) 137 138 # Update the job's dependencies to include the metahost. 139 for metahost_label in metahosts: 140 dep = models.Label.objects.get(id=metahost_label) 141 job.dependency_labels.add(dep) 142 143 for host_id in hosts: 144 models.HostQueueEntry.objects.create(job=job, host_id=host_id, 145 status=status) 146 models.IneligibleHostQueue.objects.create(job=job, host_id=host_id) 147 for label_id in metahosts: 148 models.HostQueueEntry.objects.create(job=job, meta_host_id=label_id, 149 status=status) 150 151 if hostless: 152 assert not (hosts or metahosts) 153 models.HostQueueEntry.objects.create(job=job, status=status) 154 return job 155 156 157 def _create_job_simple(self, hosts, use_metahost=False, 158 priority=0, active=False, drone_set=None, 159 parent_job_id=None): 160 """An alternative interface to _create_job""" 161 args = {'hosts' : [], 'metahosts' : []} 162 if use_metahost: 163 args['metahosts'] = hosts 164 else: 165 args['hosts'] = hosts 166 return self._create_job( 167 priority=priority, active=active, drone_set=drone_set, 168 parent_job_id=parent_job_id, **args) 169