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