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.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