#!/usr/bin/python
|
#pylint: disable-msg=C0111
|
|
# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
|
# Use of this source code is governed by a BSD-style license that can be
|
# found in the LICENSE file.
|
|
import abc
|
import os
|
|
import common
|
|
from autotest_lib.database import database_connection
|
from autotest_lib.frontend import setup_django_environment
|
from autotest_lib.frontend.afe import frontend_test_utils
|
from autotest_lib.frontend.afe import models
|
from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models
|
from autotest_lib.scheduler import monitor_db
|
from autotest_lib.scheduler import query_managers
|
from autotest_lib.scheduler import scheduler_lib
|
from autotest_lib.scheduler import scheduler_models
|
from autotest_lib.scheduler import rdb_hosts
|
from autotest_lib.scheduler import rdb_requests
|
from autotest_lib.server.cros import provision
|
|
|
# Set for verbose table creation output.
|
_DEBUG = False
|
DEFAULT_ACLS = ['Everyone', 'my_acl']
|
DEFAULT_DEPS = ['a', 'b']
|
DEFAULT_USER = 'system'
|
|
|
def get_default_job_params():
|
return {'deps': DEFAULT_DEPS, 'user': DEFAULT_USER, 'acls': DEFAULT_ACLS,
|
'priority': 0, 'parent_job_id': 0}
|
|
|
def get_default_host_params():
|
return {'deps': DEFAULT_DEPS, 'acls': DEFAULT_ACLS}
|
|
|
class FakeHost(rdb_hosts.RDBHost):
|
"""Fake host to use in unittests."""
|
|
def __init__(self, hostname, host_id, **kwargs):
|
kwargs.update({'hostname': hostname, 'id': host_id})
|
kwargs = rdb_models.AbstractHostModel.provide_default_values(
|
kwargs)
|
super(FakeHost, self).__init__(**kwargs)
|
|
|
def wire_format_response_map(response_map):
|
wire_formatted_map = {}
|
for request, response in response_map.iteritems():
|
wire_formatted_map[request] = [reply.wire_format()
|
for reply in response]
|
return wire_formatted_map
|
|
|
class DBHelper(object):
|
"""Utility class for updating the database."""
|
|
def __init__(self):
|
"""Initialized django so it uses an in memory SQLite database."""
|
self.database = (
|
database_connection.TranslatingDatabase.get_test_database(
|
translators=scheduler_lib._DB_TRANSLATORS))
|
self.database.connect(db_type='django')
|
self.database.debug = _DEBUG
|
|
|
@classmethod
|
def get_labels(cls, **kwargs):
|
"""Get a label queryset based on the kwargs."""
|
return models.Label.objects.filter(**kwargs)
|
|
|
@classmethod
|
def get_acls(cls, **kwargs):
|
"""Get an aclgroup queryset based on the kwargs."""
|
return models.AclGroup.objects.filter(**kwargs)
|
|
|
@classmethod
|
def get_host(cls, **kwargs):
|
"""Get a host queryset based on the kwargs."""
|
return models.Host.objects.filter(**kwargs)
|
|
|
@classmethod
|
def get_hqes(cls, **kwargs):
|
return models.HostQueueEntry.objects.filter(**kwargs)
|
|
|
@classmethod
|
def get_tasks(cls, **kwargs):
|
return models.SpecialTask.objects.filter(**kwargs)
|
|
|
@classmethod
|
def get_shard(cls, **kwargs):
|
return models.Shard.objects.filter(**kwargs)
|
|
|
@classmethod
|
def create_label(cls, name, **kwargs):
|
label = cls.get_labels(name=name, **kwargs)
|
return (models.Label.add_object(name=name, **kwargs)
|
if not label else label[0])
|
|
|
@classmethod
|
def create_user(cls, name):
|
user = models.User.objects.filter(login=name)
|
return models.User.add_object(login=name) if not user else user[0]
|
|
|
@classmethod
|
def create_special_task(cls, job_id=None, host_id=None,
|
task=models.SpecialTask.Task.VERIFY,
|
user='autotest-system'):
|
if job_id:
|
queue_entry = cls.get_hqes(job_id=job_id)[0]
|
host_id = queue_entry.host.id
|
else:
|
queue_entry = None
|
host = models.Host.objects.get(id=host_id)
|
owner = cls.create_user(user)
|
if not host:
|
raise ValueError('Require a host to create special tasks.')
|
return models.SpecialTask.objects.create(
|
host=host, queue_entry=queue_entry, task=task,
|
requested_by_id=owner.id)
|
|
|
@classmethod
|
def create_shard(cls, shard_hostname):
|
"""Create a shard with the given hostname if one doesn't already exist.
|
|
@param shard_hostname: The hostname of the shard.
|
"""
|
shard = cls.get_shard(hostname=shard_hostname)
|
return (models.Shard.objects.create(hostname=shard_hostname)
|
if not shard else shard[0])
|
|
|
@classmethod
|
def add_labels_to_host(cls, host, label_names=set([])):
|
label_objects = set([])
|
for label in label_names:
|
label_objects.add(cls.create_label(label))
|
host.labels.add(*label_objects)
|
|
|
@classmethod
|
def create_acl_group(cls, name):
|
aclgroup = cls.get_acls(name=name)
|
return (models.AclGroup.add_object(name=name)
|
if not aclgroup else aclgroup[0])
|
|
|
@classmethod
|
def add_deps_to_job(cls, job, dep_names=set([])):
|
label_objects = set([])
|
for label in dep_names:
|
label_objects.add(cls.create_label(label))
|
job.dependency_labels.add(*label_objects)
|
|
|
@classmethod
|
def assign_job_to_shard(cls, job_id, shard_hostname):
|
"""Assign a job to a shard.
|
|
@param job: A job object without a shard.
|
@param shard_hostname: The hostname of a shard to assign the job.
|
|
@raises ValueError: If the job already has a shard.
|
"""
|
job_filter = models.Job.objects.filter(id=job_id, shard__isnull=True)
|
if len(job_filter) != 1:
|
raise ValueError('Failed to assign job %s to shard %s' %
|
job_filter, shard_hostname)
|
job_filter.update(shard=cls.create_shard(shard_hostname))
|
|
|
@classmethod
|
def add_host_to_aclgroup(cls, host, aclgroup_names=set([])):
|
for group_name in aclgroup_names:
|
aclgroup = cls.create_acl_group(group_name)
|
aclgroup.hosts.add(host)
|
|
|
@classmethod
|
def add_user_to_aclgroups(cls, username, aclgroup_names=set([])):
|
user = cls.create_user(username)
|
for group_name in aclgroup_names:
|
aclgroup = cls.create_acl_group(group_name)
|
aclgroup.users.add(user)
|
|
|
@classmethod
|
def create_host(cls, name, deps=set([]), acls=set([]), status='Ready',
|
locked=0, lock_reason='', leased=0, protection=0, dirty=0):
|
"""Create a host.
|
|
Also adds the appropriate labels to the host, and adds the host to the
|
required acl groups.
|
|
@param name: The hostname.
|
@param kwargs:
|
deps: The labels on the host that match job deps.
|
acls: The aclgroups this host must be a part of.
|
status: The status of the host.
|
locked: 1 if the host is locked.
|
lock_reason: non-empty string if the host is locked.
|
leased: 1 if the host is leased.
|
protection: Any protection level, such as Do Not Verify.
|
dirty: 1 if the host requires cleanup.
|
|
@return: The host object for the new host.
|
"""
|
# TODO: Modify this to use the create host request once
|
# crbug.com/350995 is fixed.
|
host = models.Host.add_object(
|
hostname=name, status=status, locked=locked,
|
lock_reason=lock_reason, leased=leased,
|
protection=protection)
|
cls.add_labels_to_host(host, label_names=deps)
|
cls.add_host_to_aclgroup(host, aclgroup_names=acls)
|
|
# Though we can return the host object above, this proves that the host
|
# actually got saved in the database. For example, this will return none
|
# if save() wasn't called on the model.Host instance.
|
return cls.get_host(hostname=name)[0]
|
|
|
@classmethod
|
def update_hqe(cls, hqe_id, **kwargs):
|
"""Update the hqe with the given kwargs.
|
|
@param hqe_id: The id of the hqe to update.
|
"""
|
models.HostQueueEntry.objects.filter(id=hqe_id).update(**kwargs)
|
|
|
@classmethod
|
def update_special_task(cls, task_id, **kwargs):
|
"""Update special tasks with the given kwargs.
|
|
@param task_id: The if of the task to update.
|
"""
|
models.SpecialTask.objects.filter(id=task_id).update(**kwargs)
|
|
|
@classmethod
|
def add_host_to_job(cls, host, job_id, activate=0):
|
"""Add a host to the hqe of a job.
|
|
@param host: An instance of the host model.
|
@param job_id: The job to which we need to add the host.
|
@param activate: If true, flip the active bit on the hqe.
|
|
@raises ValueError: If the hqe for the job already has a host,
|
or if the host argument isn't a Host instance.
|
"""
|
hqe = models.HostQueueEntry.objects.get(job_id=job_id)
|
if hqe.host:
|
raise ValueError('HQE for job %s already has a host' % job_id)
|
hqe.host = host
|
hqe.save()
|
if activate:
|
cls.update_hqe(hqe.id, active=True)
|
|
|
@classmethod
|
def increment_priority(cls, job_id):
|
job = models.Job.objects.get(id=job_id)
|
job.priority = job.priority + 1
|
job.save()
|
|
|
class FileDatabaseHelper(object):
|
"""A helper class to setup a SQLite database backed by a file.
|
|
Note that initializing a file database takes significantly longer than an
|
in-memory database and should only be used for functional tests.
|
"""
|
|
DB_FILE = os.path.join(common.autotest_dir, 'host_scheduler_db')
|
|
def initialize_database_for_testing(self, db_file_path=None):
|
"""Initialize a SQLite database for testing.
|
|
To force monitor_db and the host_scheduler to use the same SQLite file
|
database, call this method before initializing the database through
|
frontend_test_utils. The host_scheduler is setup to look for the
|
host_scheduler_db when invoked with --testing.
|
|
@param db_file_path: The name of the file to use to create
|
a SQLite database. Since this database is shared across different
|
processes using a file is closer to the real world.
|
"""
|
if not db_file_path:
|
db_file_path = self.DB_FILE
|
# TODO: Move the translating database elsewhere. Monitor_db circular
|
# imports host_scheduler.
|
from autotest_lib.frontend import setup_test_environment
|
from django.conf import settings
|
self.old_django_db_name = settings.DATABASES['default']['NAME']
|
settings.DATABASES['default']['NAME'] = db_file_path
|
self.db_file_path = db_file_path
|
_db_manager = scheduler_lib.ConnectionManager(autocommit=False)
|
_db_manager.db_connection = (
|
database_connection.TranslatingDatabase.get_test_database(
|
translators=scheduler_lib._DB_TRANSLATORS))
|
|
|
def teardown_file_database(self):
|
"""Teardown django database settings."""
|
# TODO: Move the translating database elsewhere. Monitor_db circular
|
# imports host_scheduler.
|
from django.conf import settings
|
settings.DATABASES['default']['NAME'] = self.old_django_db_name
|
try:
|
os.remove(self.db_file_path)
|
except (OSError, AttributeError):
|
pass
|
|
|
class AbstractBaseRDBTester(frontend_test_utils.FrontendTestMixin):
|
|
__meta__ = abc.ABCMeta
|
_config_section = 'AUTOTEST_WEB'
|
|
|
@staticmethod
|
def get_request(dep_names, acl_names, priority=0, parent_job_id=0):
|
deps = [dep.id for dep in DBHelper.get_labels(name__in=dep_names)]
|
acls = [acl.id for acl in DBHelper.get_acls(name__in=acl_names)]
|
return rdb_requests.AcquireHostRequest(
|
deps=deps, acls=acls, host_id=None, priority=priority,
|
parent_job_id=parent_job_id)._request
|
|
|
def _release_unused_hosts(self):
|
"""Release all hosts unused by an active hqe. """
|
self.host_scheduler.tick()
|
|
|
def setUp(self, inline_host_acquisition=True, setup_tables=True):
|
"""Common setup module for tests that need a jobs/host database.
|
|
@param inline_host_acquisition: If True, the dispatcher tries to acquire
|
hosts inline with the rest of the tick.
|
"""
|
self.db_helper = DBHelper()
|
self._database = self.db_helper.database
|
# Runs syncdb setting up initial database conditions
|
self._frontend_common_setup(setup_tables=setup_tables)
|
connection_manager = scheduler_lib.ConnectionManager(autocommit=False)
|
self.god.stub_with(connection_manager, 'db_connection', self._database)
|
self.god.stub_with(monitor_db, '_db_manager', connection_manager)
|
self.god.stub_with(scheduler_models, '_db', self._database)
|
self.god.stub_with(monitor_db, '_inline_host_acquisition',
|
inline_host_acquisition)
|
self._dispatcher = monitor_db.Dispatcher()
|
self.host_scheduler = self._dispatcher._host_scheduler
|
self.host_query_manager = query_managers.AFEHostQueryManager()
|
self.job_query_manager = self._dispatcher._job_query_manager
|
self._release_unused_hosts()
|
|
|
def tearDown(self):
|
self.god.unstub_all()
|
self._database.disconnect()
|
self._frontend_common_teardown()
|
|
|
def create_job(self, user='autotest_system',
|
deps=set([]), acls=set([]), hostless_job=False,
|
priority=0, parent_job_id=None, shard_hostname=None):
|
"""Create a job owned by user, with the deps and acls specified.
|
|
This method is a wrapper around frontend_test_utils.create_job, that
|
also takes care of creating the appropriate deps for a job, and the
|
appropriate acls for the given user.
|
|
@raises ValueError: If no deps are specified for a job, since all jobs
|
need at least the metahost.
|
@raises AssertionError: If no hqe was created for the job.
|
|
@return: An instance of the job model associated with the new job.
|
"""
|
# This is a slight hack around the implementation of
|
# scheduler_models.is_hostless_job, even though a metahost is just
|
# another label to the rdb.
|
if not deps:
|
raise ValueError('Need at least one dep for metahost')
|
|
# TODO: This is a hack around the fact that frontend_test_utils still
|
# need a metahost, but metahost is treated like any other label.
|
metahost = self.db_helper.create_label(list(deps)[0])
|
job = self._create_job(metahosts=[metahost.id], priority=priority,
|
owner=user, parent_job_id=parent_job_id)
|
self.assert_(len(job.hostqueueentry_set.all()) == 1)
|
|
self.db_helper.add_deps_to_job(job, dep_names=list(deps)[1:])
|
self.db_helper.add_user_to_aclgroups(user, aclgroup_names=acls)
|
if shard_hostname:
|
self.db_helper.assign_job_to_shard(job.id, shard_hostname)
|
return models.Job.objects.filter(id=job.id)[0]
|
|
|
def assert_host_db_status(self, host_id):
|
"""Assert host state right after acquisition.
|
|
Call this method to check the status of any host leased by the
|
rdb before it has been assigned to an hqe. It must be leased and
|
ready at this point in time.
|
|
@param host_id: Id of the host to check.
|
|
@raises AssertionError: If the host is either not leased or Ready.
|
"""
|
host = models.Host.objects.get(id=host_id)
|
self.assert_(host.leased)
|
self.assert_(host.status == 'Ready')
|
|
|
def check_hosts(self, host_iter):
|
"""Sanity check all hosts in the host_gen.
|
|
@param host_iter: A generator/iterator of RDBClientHostWrappers.
|
eg: The generator returned by rdb_lib.acquire_hosts. If a request
|
was not satisfied this iterator can contain None.
|
|
@raises AssertionError: If any of the sanity checks fail.
|
"""
|
for host in host_iter:
|
if host:
|
self.assert_host_db_status(host.id)
|
self.assert_(host.leased == 1)
|
|
|
def create_suite(self, user='autotest_system', num=2, priority=0,
|
board='z', build='x', acls=set()):
|
"""Create num jobs with the same parent_job_id, board, build, priority.
|
|
@return: A dictionary with the parent job object keyed as 'parent_job'
|
and all other jobs keyed at an index from 0-num.
|
"""
|
jobs = {}
|
# Create a hostless parent job without an hqe or deps. Since the
|
# hostless job does nothing, we need to hand craft cros-version.
|
parent_job = self._create_job(owner=user, priority=priority)
|
jobs['parent_job'] = parent_job
|
build = '%s:%s' % (provision.CROS_VERSION_PREFIX, build)
|
for job_index in range(0, num):
|
jobs[job_index] = self.create_job(user=user, priority=priority,
|
deps=set([board, build]),
|
acls=acls,
|
parent_job_id=parent_job.id)
|
return jobs
|
|
|
def check_host_assignment(self, job_id, host_id):
|
"""Check is a job<->host assignment is valid.
|
|
Uses the deps of a job and the aclgroups the owner of the job is
|
in to see if the given host can be used to run the given job. Also
|
checks that the host-job assignment has Not been made, but that the
|
host is no longer in the available hosts pool.
|
|
Use this method to check host assignements made by the rdb, Before
|
they're handed off to the scheduler, since the scheduler.
|
|
@param job_id: The id of the job to use in the compatibility check.
|
@param host_id: The id of the host to check for compatibility.
|
|
@raises AssertionError: If the job and the host are incompatible.
|
"""
|
job = models.Job.objects.get(id=job_id)
|
host = models.Host.objects.get(id=host_id)
|
hqe = job.hostqueueentry_set.all()[0]
|
|
# Confirm that the host has not been assigned, either to another hqe
|
# or the this one.
|
all_hqes = models.HostQueueEntry.objects.filter(
|
host_id=host_id, complete=0)
|
self.assert_(len(all_hqes) <= 1)
|
self.assert_(hqe.host_id == None)
|
self.assert_host_db_status(host_id)
|
|
# Assert that all deps of the job are satisfied.
|
job_deps = set([d.name for d in job.dependency_labels.all()])
|
host_labels = set([l.name for l in host.labels.all()])
|
self.assert_(job_deps.intersection(host_labels) == job_deps)
|
|
# Assert that the owner of the job is in at least one of the
|
# groups that owns the host.
|
job_owner_aclgroups = set([job_acl.name for job_acl
|
in job.user().aclgroup_set.all()])
|
host_aclgroups = set([host_acl.name for host_acl
|
in host.aclgroup_set.all()])
|
self.assert_(job_owner_aclgroups.intersection(host_aclgroups))
|
